dataframe 1.3.0.0 → 2.0.0.0
raw patch · 73 files changed
+224/−23301 lines, 73 filesdep +dataframe-coredep +dataframe-csvdep +dataframe-csv-thdep −Globdep −arraydep −asyncdep ~aesondep ~dataframedep ~processPVP ok
version bump matches the API change (PVP)
Dependencies added: dataframe-core, dataframe-csv, dataframe-csv-th, dataframe-json, dataframe-lazy, dataframe-learn, dataframe-operations, dataframe-parquet, dataframe-parquet-th, dataframe-parsing, dataframe-th, dataframe-viz
Dependencies removed: Glob, array, async, attoparsec, bytestring-lexing, cassava, deepseq, granite, hashable, http-conduit, pinch, regex-tdfa, scientific, snappy-hs, stm, streamly-bytestring, streamly-core, template-haskell, temporary, these, unordered-containers, vector-algorithms, zlib, zstd
Dependency ranges changed: aeson, dataframe, process, random
API changes (from Hackage documentation)
Files
- dataframe.cabal +105/−110
- ffi/DataFrame/IO/Arrow.hs +1/−2
- src/DataFrame.hs +63/−1
- src/DataFrame/DecisionTree.hs +0/−1003
- src/DataFrame/Display.hs +0/−30
- src/DataFrame/Display/Terminal/Colours.hs +0/−14
- src/DataFrame/Display/Terminal/Plot.hs +0/−648
- src/DataFrame/Display/Terminal/PrettyPrint.hs +0/−104
- src/DataFrame/Display/Web/Plot.hs +0/−1080
- src/DataFrame/Errors.hs +0/−188
- src/DataFrame/Functions.hs +0/−676
- src/DataFrame/IO/CSV.hs +0/−794
- src/DataFrame/IO/JSON.hs +0/−133
- src/DataFrame/IO/Parquet.hs +0/−735
- src/DataFrame/IO/Parquet/Binary.hs +0/−141
- src/DataFrame/IO/Parquet/Decompress.hs +0/−32
- src/DataFrame/IO/Parquet/Dictionary.hs +0/−164
- src/DataFrame/IO/Parquet/Encoding.hs +0/−135
- src/DataFrame/IO/Parquet/Levels.hs +0/−210
- src/DataFrame/IO/Parquet/Page.hs +0/−352
- src/DataFrame/IO/Parquet/Schema.hs +0/−86
- src/DataFrame/IO/Parquet/Seeking.hs +0/−159
- src/DataFrame/IO/Parquet/Thrift.hs +0/−584
- src/DataFrame/IO/Parquet/Time.hs +0/−67
- src/DataFrame/IO/Parquet/Utils.hs +0/−384
- src/DataFrame/IO/Utils/RandomAccess.hs +0/−78
- src/DataFrame/Internal/Binary.hs +0/−94
- src/DataFrame/Internal/Column.hs +0/−1761
- src/DataFrame/Internal/DataFrame.hs +0/−335
- src/DataFrame/Internal/Expression.hs +0/−397
- src/DataFrame/Internal/Grouping.hs +0/−192
- src/DataFrame/Internal/Interpreter.hs +0/−1064
- src/DataFrame/Internal/Nullable.hs +0/−500
- src/DataFrame/Internal/Parsing.hs +0/−219
- src/DataFrame/Internal/Row.hs +0/−175
- src/DataFrame/Internal/Schema.hs +0/−242
- src/DataFrame/Internal/Statistics.hs +0/−285
- src/DataFrame/Internal/Types.hs +0/−153
- src/DataFrame/Lazy.hs +0/−3
- src/DataFrame/Lazy/IO/Binary.hs +0/−413
- src/DataFrame/Lazy/IO/CSV.hs +0/−469
- src/DataFrame/Lazy/Internal/DataFrame.hs +0/−148
- src/DataFrame/Lazy/Internal/Executor.hs +0/−656
- src/DataFrame/Lazy/Internal/LogicalPlan.hs +0/−49
- src/DataFrame/Lazy/Internal/Optimizer.hs +0/−209
- src/DataFrame/Lazy/Internal/PhysicalPlan.hs +0/−36
- src/DataFrame/Monad.hs +0/−99
- src/DataFrame/Operations/Aggregation.hs +0/−160
- src/DataFrame/Operations/Core.hs +0/−977
- src/DataFrame/Operations/Join.hs +0/−1102
- src/DataFrame/Operations/Merge.hs +0/−73
- src/DataFrame/Operations/Permutation.hs +0/−164
- src/DataFrame/Operations/Statistics.hs +0/−397
- src/DataFrame/Operations/Subset.hs +0/−540
- src/DataFrame/Operations/Transformations.hs +0/−244
- src/DataFrame/Operations/Typing.hs +0/−470
- src/DataFrame/Operators.hs +0/−329
- src/DataFrame/Synthesis.hs +0/−483
- src/DataFrame/TH.hs +24/−202
- src/DataFrame/Typed.hs +9/−0
- src/DataFrame/Typed/Access.hs +0/−55
- src/DataFrame/Typed/Aggregate.hs +0/−118
- src/DataFrame/Typed/Expr.hs +0/−644
- src/DataFrame/Typed/Freeze.hs +0/−98
- src/DataFrame/Typed/Generic.hs +0/−205
- src/DataFrame/Typed/Join.hs +0/−72
- src/DataFrame/Typed/Lazy.hs +0/−207
- src/DataFrame/Typed/Operations.hs +0/−378
- src/DataFrame/Typed/Record.hs +0/−105
- src/DataFrame/Typed/Schema.hs +0/−441
- src/DataFrame/Typed/TH.hs +21/−313
- src/DataFrame/Typed/Types.hs +0/−114
- tests/Operations/Join.hs +1/−1
dataframe.cabal view
@@ -1,6 +1,6 @@ cabal-version: 3.0 name: dataframe-version: 1.3.0.0+version: 2.0.0.0 synopsis: A fast, safe, and intuitive DataFrame library. @@ -41,115 +41,75 @@ ghc-options: -Wall +flag no-csv+ default: False+ manual: True+ description: Exclude the CSV reader/writer (@dataframe-csv@). Enable+ with @-f +no-csv@ (or @flags: no-csv@ in cabal.project)+ to trim the dep set of the meta package when CSV is not+ needed.++flag no-parquet+ default: False+ manual: True+ description: Exclude the Parquet reader/writer (@dataframe-parquet@+ plus pinch, zstd, snappy, streamly, http-conduit).+ Enable with @-f +no-parquet@.++flag no-th+ default: False+ manual: True+ description: Exclude the Template Haskell splices (@dataframe-th@,+ @dataframe-csv-th@, @dataframe-parquet-th@). Enable+ with @-f +no-th@ to drop the @template-haskell@ dep for+ downstream packages that don't use compile-time schema+ derivation.+ library import: warnings default-extensions: Strict+ -- Lock the unused-packages gate locally; haskell-ci already runs the+ -- same Werror against this stanza. Catches the case where a satellite+ -- gets added to build-depends but never re-exported.+ ghc-options: -Werror=unused-packages exposed-modules: DataFrame,- DataFrame.Lazy,- DataFrame.Functions,- DataFrame.Synthesis,- DataFrame.Display.Web.Plot,- DataFrame.Internal.Types,- DataFrame.Internal.Expression,- DataFrame.Internal.Grouping,- DataFrame.Internal.Interpreter,- DataFrame.Internal.Nullable,- DataFrame.Internal.Parsing,- DataFrame.Internal.Column,- DataFrame.Internal.Binary,- DataFrame.Internal.Statistics,- DataFrame.Display.Terminal.PrettyPrint,- DataFrame.Display.Terminal.Colours,- DataFrame.Internal.DataFrame,- DataFrame.Internal.Row,- DataFrame.Internal.Schema,- DataFrame.Errors,- DataFrame.Operations.Core,- DataFrame.Operations.Join,- DataFrame.Operations.Merge,- DataFrame.Operators,- DataFrame.Operations.Permutation,- DataFrame.Operations.Subset,- DataFrame.Operations.Statistics,- DataFrame.Operations.Transformations,- DataFrame.Operations.Typing,- DataFrame.Operations.Aggregation,- DataFrame.Display,- 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.Decompress,- DataFrame.IO.Parquet.Encoding,- DataFrame.IO.Parquet.Page,- DataFrame.IO.Parquet.Utils,- DataFrame.IO.Parquet.Seeking,- DataFrame.IO.Parquet.Time,- DataFrame.Lazy.IO.CSV,- DataFrame.Lazy.IO.Binary,- DataFrame.IO.Parquet.Schema,- DataFrame.Lazy.Internal.DataFrame,- DataFrame.Lazy.Internal.LogicalPlan,- DataFrame.Lazy.Internal.PhysicalPlan,- DataFrame.Lazy.Internal.Optimizer,- DataFrame.Lazy.Internal.Executor,- DataFrame.Monad,- DataFrame.DecisionTree,- DataFrame.Typed.Types,- DataFrame.Typed.Schema,- DataFrame.Typed.Freeze,- DataFrame.Typed.Access,- DataFrame.Typed.Operations,- DataFrame.Typed.Join,- DataFrame.Typed.Aggregate,- DataFrame.Typed.TH,- DataFrame.Typed.Expr,- DataFrame.Typed.Lazy,- DataFrame.Typed.Record,- DataFrame.Typed.Generic,- DataFrame.Typed,- DataFrame.TH+ DataFrame.Typed build-depends: base >= 4 && <5,- async >= 2.2 && < 3,- deepseq >= 1 && < 2,- aeson >= 0.11.0.0 && < 3,- array >= 0.5.4.0 && < 0.6,- temporary >= 1.3 && < 2,- attoparsec >= 0.12 && < 0.15,- bytestring >= 0.11 && < 0.13,- bytestring-lexing >= 0.5 && < 0.6,- cassava >= 0.1 && < 1,- containers >= 0.6.7 && < 0.9,- directory >= 1.3.0.0 && < 2,- granite ^>= 0.4,- hashable >= 1.2 && < 2,- process ^>= 1.6,- snappy-hs ^>= 0.1,- random >= 1.2 && < 2,- regex-tdfa >= 1.3.0 && < 2,- scientific >=0.3.1 && <0.4,- template-haskell >= 2.0 && < 3,- text >= 2.0 && < 3,- these >= 1.1 && < 2,- time >= 1.12 && < 2,- unordered-containers >= 0.1 && < 1,- vector ^>= 0.13,- vector-algorithms ^>= 0.9,- zlib >= 0.5 && < 1,- zstd >= 0.1.2.0 && < 0.2,- stm >= 2.5 && < 3,- filepath >= 1.4 && < 2,- Glob >= 0.10 && < 1,- http-conduit >= 2.3 && < 3,- pinch >= 0.5 && < 1,- streamly-core >= 0.2.3 && < 0.4,- streamly-bytestring >= 0.2.0 && < 0.4+ dataframe-core ^>= 1.0,+ dataframe-json ^>= 1.0,+ dataframe-operations ^>= 1.0,+ dataframe-parsing ^>= 1.0,+ dataframe-viz ^>= 1.0,+ dataframe-learn ^>= 1.0 + if !flag(no-csv)+ build-depends: dataframe-csv ^>= 1.0+ cpp-options: -DWITH_CSV++ if !flag(no-parquet)+ build-depends: dataframe-parquet ^>= 1.0+ cpp-options: -DWITH_PARQUET++ -- The lazy executor calls both CSV and Parquet readers directly, so+ -- it can only be pulled in when both backends are present.+ if !flag(no-csv) && !flag(no-parquet)+ build-depends: dataframe-lazy ^>= 1.0+ cpp-options: -DWITH_LAZY++ if !flag(no-th)+ build-depends: dataframe-th ^>= 1.0+ cpp-options: -DWITH_TH+ exposed-modules: DataFrame.TH,+ DataFrame.Typed.TH++ if !flag(no-th) && !flag(no-csv)+ build-depends: dataframe-csv-th ^>= 1.0+ cpp-options: -DWITH_CSV_TH++ if !flag(no-th) && !flag(no-parquet)+ build-depends: dataframe-parquet-th ^>= 1.0+ cpp-options: -DWITH_PARQUET_TH+ hs-source-dirs: src default-language: Haskell2010 @@ -160,9 +120,20 @@ exposed-modules: DataFrame.IO.Arrow DataFrame.IR DataFrame.IR.ExprJson+ -- The Arrow bridge IR loads CSV + Parquet readers, so it requires+ -- both backends. If either flag is off, skip building it.+ if flag(no-csv) || flag(no-parquet)+ buildable: False build-depends: base >= 4 && < 5, dataframe,+ dataframe-core ^>= 1.0,+ dataframe-csv ^>= 1.0,+ dataframe-json ^>= 1.0,+ dataframe-lazy ^>= 1.0,+ dataframe-operations ^>= 1.0,+ dataframe-parquet ^>= 1.0,+ dataframe-parsing ^>= 1.0, text >= 2.0 && < 3, aeson >= 0.11 && < 3, bytestring >= 0.11 && < 0.13,@@ -175,7 +146,8 @@ import: warnings main-is: Benchmark.hs build-depends: base >= 4 && < 5,- dataframe >= 1 && < 2,+ dataframe >= 1 && < 3,+ dataframe-operations ^>= 1.0, random >= 1 && < 2, time >= 1.12 && < 2, vector ^>= 0.13,@@ -187,7 +159,10 @@ import: warnings main-is: Synthesis.hs build-depends: base >= 4 && < 5,- dataframe >= 1 && < 2,+ dataframe >= 1 && < 3,+ dataframe-core ^>= 1.0,+ dataframe-learn ^>= 1.0,+ dataframe-operations ^>= 1.0, random >= 1 && < 2, text >= 2.0 && < 3 hs-source-dirs: app@@ -211,10 +186,18 @@ executable lazy-bench import: warnings main-is: LazyBenchmark.hs+ -- lazy-bench drives the lazy executor against CSV/Parquet sources, so+ -- it can only build when both backends are present.+ if flag(no-csv) || flag(no-parquet)+ buildable: False build-depends: base >= 4 && < 5, bytestring >= 0.11 && < 0.13, containers >= 0.6.7 && < 0.9,- dataframe >= 1 && < 2,+ dataframe >= 1 && < 3,+ dataframe-core ^>= 1.0,+ dataframe-lazy ^>= 1.0,+ dataframe-operations ^>= 1.0,+ dataframe-parsing ^>= 1.0, directory >= 1.3.0.0 && < 2, random >= 1 && < 2, text >= 2.0 && < 3,@@ -231,7 +214,7 @@ build-depends: base >= 4 && < 5, criterion >= 1 && < 2, process >= 1.6 && < 2,- dataframe >= 1 && < 2,+ dataframe >= 1 && < 3, random >= 1 && < 2, default-language: Haskell2010 ghc-options:@@ -243,6 +226,10 @@ import: warnings type: exitcode-stdio-1.0 main-is: Main.hs+ -- The test runner imports CSV, JSON, Parquet, Lazy, and TH-derived+ -- schema modules. All features must be enabled for it to compile.+ if flag(no-csv) || flag(no-parquet) || flag(no-th)+ buildable: False other-modules: Assertions, DecisionTree, Functions,@@ -278,13 +265,21 @@ Monad build-depends: base >= 4 && < 5, bytestring >= 0.11 && < 0.13,- dataframe >= 1 && < 2,+ dataframe >= 1 && < 3,+ dataframe-core ^>= 1.0,+ dataframe-csv ^>= 1.0,+ dataframe-json ^>= 1.0,+ dataframe-lazy ^>= 1.0,+ dataframe-learn ^>= 1.0,+ dataframe-operations ^>= 1.0,+ dataframe-parquet ^>= 1.0,+ dataframe-parsing ^>= 1.0,+ dataframe-th ^>= 1.0, HUnit ^>= 1.6, QuickCheck >= 2 && < 3, random-shuffle >= 0.0.4 && < 1, random >= 1 && < 2, text >= 2.0 && < 3,- these >= 1.1 && < 2, time >= 1.12 && < 2, vector ^>= 0.13, containers >= 0.6.7 && < 0.9
ffi/DataFrame/IO/Arrow.hs view
@@ -30,8 +30,7 @@ import Type.Reflection (typeRep) import DataFrame.Internal.Column (Column (..))-import DataFrame.Internal.DataFrame (DataFrame (..))-import DataFrame.Operations.Core (fromNamedColumns)+import DataFrame.Internal.DataFrame (DataFrame (..), fromNamedColumns) -- --------------------------------------------------------------------------- -- Opaque phantom types for the Arrow structs
src/DataFrame.hs view
@@ -1,3 +1,5 @@+{-# LANGUAGE CPP #-}+ {- | Module : DataFrame Copyright : (c) 2025@@ -213,11 +215,28 @@ -- * Types module Schema,+#ifdef WITH_TH+ module SchemaTH,+#endif -- * I/O+#ifdef WITH_CSV module CSV,+#endif+#ifdef WITH_PARQUET module Parquet,+#endif+ module JSON, + -- * Lazy query engine+#ifdef WITH_LAZY+ module Lazy,+#endif++ -- * Feature synthesis & decision trees+ module Synthesis,+ module DecisionTree,+ -- * Type conversion module Typing, @@ -237,13 +256,20 @@ module Record, -- * Template Haskell column-binding splices+#ifdef WITH_TH module TH,+#endif -- * Plotting module Plot, ) where +-- DecisionTree defines its own `percentile`/`percentiles` helpers; the+-- public versions surfaced through the meta module are the ones from+-- Operations.Statistics / DataFrame.Synthesis. Hide the duplicates to+-- avoid ambiguous-export errors.+import DataFrame.DecisionTree as DecisionTree hiding (percentile, percentiles) import DataFrame.Display as Display ( DisplayOptions (..), defaultDisplayOptions,@@ -251,6 +277,7 @@ ) import DataFrame.Display.Terminal.Plot as Plot import DataFrame.Errors as Errors+#ifdef WITH_CSV import DataFrame.IO.CSV as CSV ( HeaderSpec (..), ReadOptions (..),@@ -266,6 +293,12 @@ writeCsv, writeSeparated, )+#endif+import DataFrame.IO.JSON as JSON (+ readJSON,+ readJSONEither,+ )+#ifdef WITH_PARQUET import DataFrame.IO.Parquet as Parquet ( ParquetReadOptions (..), defaultParquetReadOptions,@@ -274,6 +307,7 @@ readParquetFilesWithOpts, readParquetWithOpts, )+#endif import DataFrame.Internal.Column as Column ( Column, fromList,@@ -290,8 +324,11 @@ DataFrame, GroupedDataFrame, TruncateConfig (..),+ columnNames, defaultTruncateConfig, empty,+ fromNamedColumns,+ insertColumn, null, toCsv, toCsv',@@ -310,10 +347,28 @@ toRowVector, ) import DataFrame.Internal.Schema as Schema (- deriveSchema, makeSchema, schemaType, )+#ifdef WITH_TH+import DataFrame.Internal.Schema.TH as SchemaTH (deriveSchema)+#endif+#ifdef WITH_LAZY+-- Re-export the lazy query engine's entry points. Only types and source+-- constructors are surfaced; the operator surface (filter/select/derive/...)+-- collides with the eager API already re-exported above, so users wanting+-- full lazy access should `import qualified DataFrame.Lazy as L`.+import DataFrame.Lazy as Lazy (+ LazyDataFrame,+ fromDataFrame,+ runDataFrame,+ scanCsv,+ scanCsvWith,+ scanParquet,+ scanSeparated,+ scanSeparatedWith,+ )+#endif import DataFrame.Operations.Aggregation as Aggregation ( aggregate, distinct,@@ -355,6 +410,7 @@ summarize, variance, )+import DataFrame.Synthesis as Synthesis import DataFrame.Operations.Subset as Subset ( SelectionCriteria, byIndexRange,@@ -405,14 +461,20 @@ parseDefaults, ) import DataFrame.Operators as Operators+#ifdef WITH_TH import DataFrame.TH as TH ( declareColumns,+#ifdef WITH_CSV_TH declareColumnsFromCsvFile, declareColumnsFromCsvWithOpts,+#endif+#ifdef WITH_PARQUET_TH declareColumnsFromParquetFile,+#endif declareColumnsWithPrefix, declareColumnsWithPrefix', )+#endif import DataFrame.Typed.Record as Record ( HasSchema (..), fromRecords,
− src/DataFrame/DecisionTree.hs
@@ -1,1003 +0,0 @@-{-# LANGUAGE AllowAmbiguousTypes #-}-{-# LANGUAGE BangPatterns #-}-{-# LANGUAGE CPP #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE RankNTypes #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}--module DataFrame.DecisionTree where--import qualified DataFrame.Functions as F-import DataFrame.Internal.Column-import DataFrame.Internal.DataFrame (DataFrame (..), unsafeGetColumn)-import DataFrame.Internal.Expression (Expr (..), eSize, eqExpr, getColumns)-import DataFrame.Internal.Interpreter (interpret)-import DataFrame.Internal.Statistics (percentileOrd')-import DataFrame.Internal.Types-import DataFrame.Operations.Core (columnNames, nRows)-import DataFrame.Operations.Subset (exclude, filterWhere)--import Control.Exception (throw)-import Control.Monad (guard)-import Data.Function (on)-#if MIN_VERSION_base(4,20,0)-import Data.List (maximumBy, minimumBy, nub, nubBy, sort, sortBy)-#else-import Data.List (foldl', maximumBy, minimumBy, nub, nubBy, sort, sortBy)-#endif-import Data.Int (Int16, Int32, Int64, Int8)-import qualified Data.Map.Strict as M-import Data.Proxy (Proxy (..))-import qualified Data.Text as T-import Data.Type.Equality-import qualified Data.Vector as V-import qualified Data.Vector.Unboxed as VU-import Data.Word (Word16, Word32, Word64, Word8)-import Type.Reflection (SomeTypeRep (..), typeRep)--import DataFrame.Operators--{- | Declares which column types support ordering for decision tree splits.--Use 'orderable' to register a type, and '<>' to combine:--@-defaultTreeConfig- { columnOrdering = defaultColumnOrdering <> orderable \@MyCustomType- }-@--}-newtype ColumnOrdering = ColumnOrdering (M.Map SomeTypeRep OrdDict)--instance Semigroup ColumnOrdering where- ColumnOrdering a <> ColumnOrdering b = ColumnOrdering (a <> b)--instance Monoid ColumnOrdering where- mempty = ColumnOrdering M.empty---- | Register a type as orderable for decision tree splits.-orderable :: forall a. (Columnable a, Ord a) => ColumnOrdering-orderable = ColumnOrdering (M.singleton (SomeTypeRep (typeRep @a)) (OrdDict (Proxy @a)))---- | All standard numeric, text, and primitive types.-defaultColumnOrdering :: ColumnOrdering-defaultColumnOrdering =- mconcat- [ orderable @Int- , orderable @Int8- , orderable @Int16- , orderable @Int32- , orderable @Int64- , orderable @Word- , orderable @Word8- , orderable @Word16- , orderable @Word32- , orderable @Word64- , orderable @Integer- , orderable @Double- , orderable @Float- , orderable @Bool- , orderable @Char- , orderable @T.Text- , orderable @String- ]---- Internal: existential Ord dictionary.-data OrdDict where- OrdDict :: (Columnable a, Ord a) => Proxy a -> OrdDict---- Internal: look up Ord for type @a@.-withOrdFrom ::- forall a r. (Columnable a) => ColumnOrdering -> ((Ord a) => r) -> Maybe r-withOrdFrom (ColumnOrdering m) k = case M.lookup (SomeTypeRep (typeRep @a)) m of- Just (OrdDict (_ :: Proxy b)) -> case testEquality (typeRep @a) (typeRep @b) of- Just Refl -> Just k- Nothing -> Nothing- Nothing -> Nothing--data TreeConfig = TreeConfig- { maxTreeDepth :: Int- , minSamplesSplit :: Int- , minLeafSize :: Int- , percentiles :: [Int]- , expressionPairs :: Int- , synthConfig :: SynthConfig- , taoIterations :: Int- , taoConvergenceTol :: Double- , columnOrdering :: ColumnOrdering- }--data SynthConfig = SynthConfig- { maxExprDepth :: Int- , boolExpansion :: Int- , disallowedCombinations :: [(T.Text, T.Text)]- , complexityPenalty :: Double- , enableStringOps :: Bool- , enableCrossCols :: Bool- , enableArithOps :: Bool- }- deriving (Eq, Show)--defaultSynthConfig :: SynthConfig-defaultSynthConfig =- SynthConfig- { maxExprDepth = 2- , boolExpansion = 2- , disallowedCombinations = []- , complexityPenalty = 0.05- , enableStringOps = True- , enableCrossCols = True- , enableArithOps = True- }--defaultTreeConfig :: TreeConfig-defaultTreeConfig =- TreeConfig- { maxTreeDepth = 4- , minSamplesSplit = 5- , minLeafSize = 1- , percentiles = [0, 10 .. 100]- , expressionPairs = 10- , synthConfig = defaultSynthConfig- , taoIterations = 10- , taoConvergenceTol = 1e-6- , columnOrdering = defaultColumnOrdering- }--data Tree a- = Leaf !a- | Branch !(Expr Bool) !(Tree a) !(Tree a)- deriving (Show)--treeDepth :: Tree a -> Int-treeDepth (Leaf _) = 0-treeDepth (Branch _ l r) = 1 + max (treeDepth l) (treeDepth r)--treeToExpr :: (Columnable a) => Tree a -> Expr a-treeToExpr (Leaf v) = Lit v-treeToExpr (Branch cond left right) =- F.ifThenElse cond (treeToExpr left) (treeToExpr right)---- | Fit a TAO decision tree-fitDecisionTree ::- forall a.- (Columnable a, Ord a) =>- TreeConfig ->- Expr a ->- DataFrame ->- Expr a-fitDecisionTree cfg (Col target) df =- let- conds =- nubBy eqExpr $- numericConditions cfg (exclude [target] df)- ++ generateConditionsOld cfg (exclude [target] df)-- initialTree = buildGreedyTree @a cfg (maxTreeDepth cfg) target conds df-- indices = V.enumFromN 0 (nRows df)-- optimizedTree = taoOptimize @a cfg target conds df indices initialTree- in- pruneExpr (treeToExpr optimizedTree)-fitDecisionTree _ expr _ = error $ "Cannot create tree for compound expression: " ++ show expr--taoOptimize ::- forall a.- (Columnable a, Ord a) =>- TreeConfig ->- T.Text -> -- Target column name- [Expr Bool] -> -- Candidate conditions- DataFrame -> -- Full dataset- V.Vector Int -> -- Indices of points reaching the root- Tree a -> -- Current tree- Tree a-taoOptimize cfg target conds df rootIndices initialTree =- go 0 initialTree (computeTreeLoss @a target df rootIndices initialTree)- where- go :: Int -> Tree a -> Double -> Tree a- go iter tree prevLoss- | iter >= taoIterations cfg = pruneDead tree- | otherwise =- let- tree' = taoIteration @a cfg target conds df rootIndices tree-- newLoss = computeTreeLoss @a target df rootIndices tree'- improvement = prevLoss - newLoss- in- if improvement < taoConvergenceTol cfg- then pruneDead tree'- else go (iter + 1) tree' newLoss--taoIteration ::- forall a.- (Columnable a, Ord a) =>- TreeConfig ->- T.Text ->- [Expr Bool] ->- DataFrame ->- V.Vector Int ->- Tree a ->- Tree a-taoIteration cfg target conds df rootIndices tree =- let depth = treeDepth tree- in foldl'- (optimizeDepthLevel @a cfg target conds df rootIndices)- tree- [depth, depth - 1 .. 0] -- Bottom to top--optimizeDepthLevel ::- forall a.- (Columnable a, Ord a) =>- TreeConfig ->- T.Text ->- [Expr Bool] ->- DataFrame ->- V.Vector Int ->- Tree a ->- Int -> -- Target depth- Tree a-optimizeDepthLevel cfg target conds df rootIndices tree = optimizeAtDepth @a cfg target conds df rootIndices tree 0--optimizeAtDepth ::- forall a.- (Columnable a, Ord a) =>- TreeConfig ->- T.Text ->- [Expr Bool] ->- DataFrame ->- V.Vector Int ->- Tree a ->- Int ->- Int ->- Tree a-optimizeAtDepth cfg target conds df indices tree currentDepth targetDepth- | currentDepth == targetDepth =- optimizeNode @a cfg target conds df indices tree- | otherwise = case tree of- Leaf v -> Leaf v- Branch cond left right ->- let- (indicesL, indicesR) = partitionIndices cond df indices- left' =- optimizeAtDepth @a- cfg- target- conds- df- indicesL- left- (currentDepth + 1)- targetDepth- right' =- optimizeAtDepth @a- cfg- target- conds- df- indicesR- right- (currentDepth + 1)- targetDepth- in- Branch cond left' right'--optimizeNode ::- forall a.- (Columnable a, Ord a) =>- TreeConfig ->- T.Text ->- [Expr Bool] ->- DataFrame ->- V.Vector Int ->- Tree a ->- Tree a-optimizeNode cfg target conds df indices tree- | V.null indices = tree- | otherwise = case tree of- Leaf _ -> Leaf (majorityValueFromIndices @a target df indices)- Branch oldCond left right ->- let- newCond = findBestSplitTAO @a cfg target conds df indices left right oldCond-- (newIndicesL, newIndicesR) = partitionIndices newCond df indices- in- if V.length newIndicesL < minLeafSize cfg- || V.length newIndicesR < minLeafSize cfg- then Leaf (majorityValueFromIndices @a target df indices)- else Branch newCond left right--findBestSplitTAO ::- forall a.- (Columnable a) =>- TreeConfig ->- T.Text ->- [Expr Bool] ->- DataFrame ->- V.Vector Int ->- Tree a -> -- Left subtree (FIXED)- Tree a -> -- Right subtree (FIXED)- Expr Bool -> -- Current condition (fallback)- Expr Bool-findBestSplitTAO cfg target conds df indices leftTree rightTree currentCond- | V.null indices = currentCond- | null validConds = currentCond- | otherwise =- let- carePoints = identifyCarePoints @a target df indices leftTree rightTree- in- if null carePoints- then currentCond- else- let- evalSplit :: Expr Bool -> Int- evalSplit cond = countCarePointErrors cond df carePoints-- evalWithPenalty c =- let errors = evalSplit c- penalty =- floor- ( complexityPenalty (synthConfig cfg)- * fromIntegral (eSize c)- )- in errors + penalty-- sortedConds =- take (expressionPairs cfg) $- sortBy (compare `on` evalWithPenalty) validConds-- expandedConds =- boolExprs- df- sortedConds- sortedConds- 0- (boolExpansion (synthConfig cfg))- in- if null expandedConds- then currentCond- else minimumBy (compare `on` evalWithPenalty) expandedConds- where- validConds = filter isValidSplit conds- isValidSplit c =- let (t, f) = partitionIndices c df indices- in V.length t >= minLeafSize cfg && V.length f >= minLeafSize cfg---- | A care point with its index and which direction leads to correct classification-data CarePoint = CarePoint- { cpIndex :: !Int- , cpCorrectDir :: !Direction -- Which child classifies this point correctly- }- deriving (Eq, Show)--data Direction = GoLeft | GoRight- deriving (Eq, Show)--{- | Identify care points: points where exactly one subtree classifies correctly-- For each point reaching the node:- 1. Compute what label the left subtree would predict- 2. Compute what label the right subtree would predict- 3. If exactly one matches the true label, it's a care point- 4. Record which direction leads to correct classification--}-identifyCarePoints ::- forall a.- (Columnable a) =>- T.Text ->- DataFrame ->- V.Vector Int ->- Tree a -> -- Left subtree- Tree a -> -- Right subtree- [CarePoint]-identifyCarePoints target df indices leftTree rightTree =- case interpret @a df (Col target) of- Left _ -> []- Right (TColumn column) ->- case toVector @a column of- Left _ -> []- Right targetVals ->- V.toList $ V.mapMaybe (checkPoint targetVals) indices- where- checkPoint :: V.Vector a -> Int -> Maybe CarePoint- checkPoint targetVals idx =- let- trueLabel = targetVals V.! idx- leftPred = predictWithTree @a target df idx leftTree- rightPred = predictWithTree @a target df idx rightTree- leftCorrect = leftPred == trueLabel- rightCorrect = rightPred == trueLabel- in- case (leftCorrect, rightCorrect) of- (True, False) -> Just $ CarePoint idx GoLeft- (False, True) -> Just $ CarePoint idx GoRight- _ -> Nothing -- Don't-care point (both correct or both wrong)---- | Predict the label for a single point using a fixed tree-predictWithTree ::- forall a.- (Columnable a) =>- T.Text ->- DataFrame ->- Int -> -- Row index- Tree a ->- a-predictWithTree _target _df _idx (Leaf v) = v-predictWithTree target df idx (Branch cond left right) =- case interpret @Bool df cond of- Left _ -> predictWithTree @a target df idx left -- Default to left on error- Right (TColumn column) ->- case toVector @Bool column of- Left _ -> predictWithTree @a target df idx left- Right boolVals ->- if boolVals V.! idx- then predictWithTree @a target df idx left- else predictWithTree @a target df idx right--countCarePointErrors :: Expr Bool -> DataFrame -> [CarePoint] -> Int-countCarePointErrors cond df carePoints =- case interpret @Bool df cond of- Left _ -> length carePoints- Right (TColumn column) ->- case toVector @Bool column of- Left _ -> length carePoints- Right boolVals ->- length $ filter (isMisclassified boolVals) carePoints- where- isMisclassified :: V.Vector Bool -> CarePoint -> Bool- isMisclassified boolVals cp =- let goesLeft = boolVals V.! cpIndex cp- shouldGoLeft = cpCorrectDir cp == GoLeft- in goesLeft /= shouldGoLeft--partitionIndices ::- Expr Bool -> DataFrame -> V.Vector Int -> (V.Vector Int, V.Vector Int)-partitionIndices cond df indices =- case interpret @Bool df cond of- Left _ -> (indices, V.empty)- Right (TColumn column) ->- case toVector @Bool column of- Left _ -> (indices, V.empty)- Right boolVals ->- V.partition (boolVals V.!) indices--majorityValueFromIndices ::- forall a.- (Columnable a, Ord a) =>- T.Text ->- DataFrame ->- V.Vector Int ->- a-majorityValueFromIndices target df indices =- case interpret @a df (Col target) of- Left e -> throw e- Right (TColumn column) ->- case toVector @a column of- Left e -> throw e- Right vals ->- let counts =- V.foldl'- (\acc i -> M.insertWith (+) (vals V.! i) (1 :: Int) acc)- M.empty- indices- in if M.null counts- then error "Empty indices in majorityValueFromIndices"- else fst $ maximumBy (compare `on` snd) (M.toList counts)--computeTreeLoss ::- forall a.- (Columnable a) =>- T.Text ->- DataFrame ->- V.Vector Int ->- Tree a ->- Double-computeTreeLoss target df indices tree- | V.null indices = 0- | otherwise =- case interpret @a df (Col target) of- Left _ -> 1.0- Right (TColumn column) ->- case toVector @a column of- Left _ -> 1.0- Right targetVals ->- let- n = V.length indices- errors =- V.length $- V.filter- (\i -> targetVals V.! i /= predictWithTree @a target df i tree)- indices- in- fromIntegral errors / fromIntegral n--pruneDead :: Tree a -> Tree a-pruneDead (Leaf v) = Leaf v-pruneDead (Branch cond left right) =- let- left' = pruneDead left- right' = pruneDead right- in- Branch cond left' right'--pruneExpr :: forall a. (Columnable a) => Expr a -> Expr a-pruneExpr (If cond trueBranch falseBranch) =- let t = pruneExpr trueBranch- f = pruneExpr falseBranch- in if eqExpr t f- then t- else case (t, f) of- (If condInner tInner _, _) | eqExpr cond condInner -> If cond tInner f- (_, If condInner _ fInner) | eqExpr cond condInner -> If cond t fInner- _ -> If cond t f-pruneExpr (Unary op e) = Unary op (pruneExpr e)-pruneExpr (Binary op l r) = Binary op (pruneExpr l) (pruneExpr r)-pruneExpr e = e--buildGreedyTree ::- forall a.- (Columnable a, Ord a) =>- TreeConfig ->- Int ->- T.Text ->- [Expr Bool] ->- DataFrame ->- Tree a-buildGreedyTree cfg depth target conds df- | depth <= 0 || nRows df <= minSamplesSplit cfg =- Leaf (majorityValue @a target df)- | otherwise =- case findBestGreedySplit @a cfg target conds df of- Nothing -> Leaf (majorityValue @a target df)- Just bestCond ->- let (dfTrue, dfFalse) = partitionDataFrame bestCond df- in if nRows dfTrue < minLeafSize cfg || nRows dfFalse < minLeafSize cfg- then Leaf (majorityValue @a target df)- else- Branch- bestCond- (buildGreedyTree @a cfg (depth - 1) target conds dfTrue)- (buildGreedyTree @a cfg (depth - 1) target conds dfFalse)--findBestGreedySplit ::- forall a.- (Columnable a, Ord a) =>- TreeConfig -> T.Text -> [Expr Bool] -> DataFrame -> Maybe (Expr Bool)-findBestGreedySplit cfg target conds df =- let- initialImpurity = calculateGini @a target df- calculateComplexity c = complexityPenalty (synthConfig cfg) * fromIntegral (eSize c)-- evalGain :: Expr Bool -> (Double, Int)- evalGain cond =- let (t, f) = partitionDataFrame cond df- n = fromIntegral @Int @Double (nRows df)- weightT = fromIntegral @Int @Double (nRows t) / n- weightF = fromIntegral @Int @Double (nRows f) / n- newImpurity =- weightT * calculateGini @a target t- + weightF * calculateGini @a target f- in ( (initialImpurity - newImpurity) - calculateComplexity cond- , negate (eSize cond)- )-- validConds =- filter- ( \c ->- let (t, f) = partitionDataFrame c df- in nRows t >= minLeafSize cfg && nRows f >= minLeafSize cfg- )- conds-- sortedConditions =- map fst $- take- (expressionPairs cfg)- ( filter- (\(c, v) -> ((> negate (calculateComplexity c)) . fst) v)- (sortBy (flip compare `on` snd) (map (\c -> (c, evalGain c)) validConds))- )- in- if null sortedConditions- then Nothing- else- Just $- maximumBy- (compare `on` evalGain)- ( boolExprs- df- sortedConditions- sortedConditions- 0- (boolExpansion (synthConfig cfg))- )---- | Unifies non-nullable and nullable Double expressions for feature generation.-data NumExpr- = NDouble !(Expr Double)- | NMaybeDouble !(Expr (Maybe Double))--numExprCols :: NumExpr -> [T.Text]-numExprCols (NDouble e) = getColumns e-numExprCols (NMaybeDouble e) = getColumns e--numExprEq :: NumExpr -> NumExpr -> Bool-numExprEq (NDouble e1) (NDouble e2) = eqExpr e1 e2-numExprEq (NMaybeDouble e1) (NMaybeDouble e2) = eqExpr e1 e2-numExprEq _ _ = False--combineNumExprs :: NumExpr -> NumExpr -> [NumExpr]-combineNumExprs (NDouble e1) (NDouble e2) =- [ NDouble (e1 .+ e2)- , NDouble (e1 .- e2)- , NDouble (e1 .* e2)- , NDouble- (F.ifThenElse (e2 ./= F.lit (0 :: Double)) (e1 ./ e2) (F.lit (0 :: Double)))- ]-combineNumExprs (NDouble e1) (NMaybeDouble e2) =- [ NMaybeDouble (e1 .+ e2)- , NMaybeDouble (e1 .- e2)- , NMaybeDouble (e1 .* e2)- , NMaybeDouble- ( F.ifThenElse- (F.fromMaybe False (e2 ./= F.lit (0 :: Double)))- (e1 ./ e2)- (F.lit (Nothing :: Maybe Double))- )- ]-combineNumExprs (NMaybeDouble e1) (NDouble e2) =- [ NMaybeDouble (e1 .+ e2)- , NMaybeDouble (e1 .- e2)- , NMaybeDouble (e1 .* e2)- , NMaybeDouble- ( F.ifThenElse- (e2 ./= F.lit (0 :: Double))- (e1 ./ e2)- (F.lit (Nothing :: Maybe Double))- )- ]-combineNumExprs (NMaybeDouble e1) (NMaybeDouble e2) =- [ NMaybeDouble (e1 .+ e2)- , NMaybeDouble (e1 .- e2)- , NMaybeDouble (e1 .* e2)- , NMaybeDouble- ( F.ifThenElse- (F.fromMaybe False (e2 ./= F.lit (0 :: Double)))- (e1 ./ e2)- (F.lit (Nothing :: Maybe Double))- )- ]--numericConditions :: TreeConfig -> DataFrame -> [Expr Bool]-numericConditions = generateNumericConds--generateNumericConds :: TreeConfig -> DataFrame -> [Expr Bool]-generateNumericConds cfg df = do- expr <- numericExprsWithTerms (synthConfig cfg) df- let thresholds = numericThresholds expr- threshold <- thresholds- numericCondsFromExpr expr threshold- where- numericThresholds (NDouble e) = map (\p -> percentile p e df) (percentiles cfg)- numericThresholds (NMaybeDouble e) = map (\p -> percentile p (F.fromMaybe 0 e) df) (percentiles cfg)-- numericCondsFromExpr (NDouble e) t =- [e .<= F.lit t, e .>= F.lit t, e .< F.lit t, e .> F.lit t]- numericCondsFromExpr (NMaybeDouble e) t =- [ F.fromMaybe False (e .<= F.lit t)- , F.fromMaybe False (e .>= F.lit t)- , F.fromMaybe False (e .< F.lit t)- , F.fromMaybe False (e .> F.lit t)- ]--numericExprsWithTerms :: SynthConfig -> DataFrame -> [NumExpr]-numericExprsWithTerms cfg df =- concatMap (numericExprs cfg df [] 0) [0 .. maxExprDepth cfg]--numericCols :: DataFrame -> [NumExpr]-numericCols df = concatMap extract (columnNames df)- where- extract colName = case unsafeGetColumn colName df of- UnboxedColumn Nothing (_ :: VU.Vector b) ->- case testEquality (typeRep @b) (typeRep @Double) of- Just Refl -> [NDouble (Col colName)]- Nothing -> case sIntegral @b of- STrue -> [NDouble (F.toDouble (Col @b colName))]- SFalse -> []- BoxedColumn (Just _) (_ :: V.Vector b) ->- case testEquality (typeRep @b) (typeRep @Double) of- Just Refl -> [NMaybeDouble (Col @(Maybe b) colName)]- Nothing -> case sIntegral @b of- STrue ->- [NMaybeDouble (F.whenPresent (realToFrac @b @Double) (Col @(Maybe b) colName))]- SFalse -> []- UnboxedColumn (Just _) (_ :: VU.Vector b) ->- case testEquality (typeRep @b) (typeRep @Double) of- Just Refl -> [NMaybeDouble (Col @(Maybe b) colName)]- Nothing -> case sIntegral @b of- STrue ->- [NMaybeDouble (F.whenPresent (realToFrac @b @Double) (Col @(Maybe b) colName))]- SFalse -> []- _ -> []--numericExprs ::- SynthConfig -> DataFrame -> [NumExpr] -> Int -> Int -> [NumExpr]-numericExprs cfg df prevExprs depth maxDepth- | depth == 0 = baseExprs ++ numericExprs cfg df baseExprs (depth + 1) maxDepth- | depth >= maxDepth = []- | otherwise =- combinedExprs ++ numericExprs cfg df combinedExprs (depth + 1) maxDepth- where- baseExprs = numericCols df- combinedExprs- | not (enableArithOps cfg) = []- | otherwise = do- e1 <- prevExprs- e2 <- baseExprs- let cols = numExprCols e1 <> numExprCols e2- guard- ( not (numExprEq e1 e2)- && not- ( any- (\(l, r) -> l `elem` cols && r `elem` cols)- (disallowedCombinations cfg)- )- )- combineNumExprs e1 e2--boolExprs ::- DataFrame -> [Expr Bool] -> [Expr Bool] -> Int -> Int -> [Expr Bool]-boolExprs df baseExprs prevExprs depth maxDepth- | depth == 0 =- baseExprs ++ boolExprs df baseExprs prevExprs (depth + 1) maxDepth- | depth >= maxDepth = []- | otherwise =- combinedExprs ++ boolExprs df baseExprs combinedExprs (depth + 1) maxDepth- where- combinedExprs = do- e1 <- prevExprs- e2 <- baseExprs- guard (Prelude.not (eqExpr e1 e2))- [F.and e1 e2, F.or e1 e2]--generateConditionsOld :: TreeConfig -> DataFrame -> [Expr Bool]-generateConditionsOld cfg df =- let- ords = columnOrdering cfg- genConds :: T.Text -> [Expr Bool]- genConds colName = case unsafeGetColumn colName df of- (BoxedColumn Nothing (column :: V.Vector a)) ->- case withOrdFrom @a ords (map (Lit . (`percentileOrd'` column)) [1, 25, 75, 99]) of- Just ps -> map (\p -> Col @a colName .==. p) ps- Nothing -> []- (BoxedColumn (Just _) (column :: V.Vector a)) -> case sFloating @a of- STrue -> [] -- handled by numericCols / numericExprs- SFalse -> case sIntegral @a of- STrue -> [] -- handled by numericCols / numericExprs- SFalse ->- case withOrdFrom @a- ords- (map (Lit . Just . (`percentileOrd'` column)) [1, 25, 75, 99]) of- Just ps -> map (\p -> Col @(Maybe a) colName .==. p) ps- Nothing -> []- (UnboxedColumn _ (_ :: VU.Vector a)) -> []-- columnConds =- concatMap- colConds- [ (l, r)- | l <- columnNames df- , r <- columnNames df- , not- ( any- (\(l', r') -> sort [l', r'] == sort [l, r])- (disallowedCombinations (synthConfig cfg))- )- ]- where- colConds (!l, !r) = case (unsafeGetColumn l df, unsafeGetColumn r df) of- ( BoxedColumn Nothing (_col1 :: V.Vector a)- , BoxedColumn Nothing (_ :: V.Vector b)- ) ->- case testEquality (typeRep @a) (typeRep @b) of- Nothing -> []- Just Refl -> [Col @a l .==. Col @a r]- (UnboxedColumn _ (_ :: VU.Vector a), UnboxedColumn _ (_ :: VU.Vector b)) -> []- ( BoxedColumn (Just _) (_ :: V.Vector a)- , BoxedColumn (Just _) (_ :: V.Vector b)- ) -> case testEquality (typeRep @a) (typeRep @b) of- Nothing -> []- Just Refl -> case testEquality (typeRep @a) (typeRep @T.Text) of- Nothing ->- case withOrdFrom @a ords [Col @(Maybe a) l .<=. Col @(Maybe a) r] of- Just leExprs ->- leExprs ++ [Col @(Maybe a) l .==. Col @(Maybe a) r]- Nothing -> [Col @(Maybe a) l .==. Col @(Maybe a) r]- Just Refl -> [Col @(Maybe a) l .==. Col @(Maybe a) r]- _ -> []- in- concatMap genConds (columnNames df) ++ columnConds--partitionDataFrame :: Expr Bool -> DataFrame -> (DataFrame, DataFrame)-partitionDataFrame cond df = (filterWhere cond df, filterWhere (F.not cond) df)--calculateGini ::- forall a. (Columnable a, Ord a) => T.Text -> DataFrame -> Double-calculateGini target df =- let n = fromIntegral $ nRows df- counts = getCounts @a target df- numClasses = fromIntegral $ M.size counts- probs = map (\c -> (fromIntegral c + 1) / (n + numClasses)) (M.elems counts)- in if n == 0 then 0 else 1 - sum (map (^ (2 :: Int)) probs)--majorityValue :: forall a. (Columnable a, Ord a) => T.Text -> DataFrame -> a-majorityValue target df =- let counts = getCounts @a target df- in if M.null counts- then error "Empty DataFrame in leaf"- else fst $ maximumBy (compare `on` snd) (M.toList counts)--getCounts ::- forall a. (Columnable a, Ord a) => T.Text -> DataFrame -> M.Map a Int-getCounts target df =- case interpret @a df (Col target) of- Left e -> throw e- Right (TColumn column) ->- case toVector @a column of- Left e -> throw e- Right vals -> foldl' (\acc x -> M.insertWith (+) x 1 acc) M.empty (V.toList vals)--percentile :: Int -> Expr Double -> DataFrame -> Double-percentile p expr df =- case interpret @Double df expr of- Left _ -> 0- Right (TColumn column) ->- case toVector @Double column of- Left _ -> 0- Right vals ->- let sorted = V.fromList $ sort $ V.toList vals- n = V.length sorted- idx = min (n - 1) $ max 0 $ (p * n) `div` 100- in if n == 0 then 0 else sorted V.! idx--buildTree ::- forall a.- (Columnable a, Ord a) =>- TreeConfig ->- Int ->- T.Text ->- [Expr Bool] ->- DataFrame ->- Expr a-buildTree cfg depth target conds df =- let- tree = buildGreedyTree @a cfg depth target conds df- indices = V.enumFromN 0 (nRows df)- optimized = taoOptimize @a cfg target conds df indices tree- in- pruneExpr (treeToExpr optimized)--findBestSplit ::- forall a.- (Columnable a, Ord a) =>- TreeConfig -> T.Text -> [Expr Bool] -> DataFrame -> Maybe (Expr Bool)-findBestSplit = findBestGreedySplit @a--pruneTree :: forall a. (Columnable a) => Expr a -> Expr a-pruneTree = pruneExpr---- | A tree where each leaf stores a class-probability distribution.-type ProbTree a = Tree (M.Map a Double)---- | Compute normalised class probabilities from a subset of training rows.-probsFromIndices ::- forall a.- (Columnable a, Ord a) =>- T.Text ->- DataFrame ->- V.Vector Int ->- M.Map a Double-probsFromIndices target df indices =- case interpret @a df (Col target) of- Left _ -> M.empty- Right (TColumn column) ->- case toVector @a column of- Left _ -> M.empty- Right vals ->- let counts =- V.foldl'- (\acc i -> M.insertWith (+) (vals V.! i) (1 :: Int) acc)- M.empty- indices- total = fromIntegral (V.length indices) :: Double- in M.map (\c -> fromIntegral c / total) counts--{- | Annotate a fitted 'Tree a' with class distributions by routing the- training data through it. The split conditions are preserved; only the- leaf values change from a majority label to a probability map.--}-buildProbTree ::- forall a.- (Columnable a, Ord a) =>- Tree a ->- T.Text ->- DataFrame ->- V.Vector Int ->- ProbTree a-buildProbTree (Leaf _) target df indices =- Leaf (probsFromIndices @a target df indices)-buildProbTree (Branch cond left right) target df indices =- let (indicesL, indicesR) = partitionIndices cond df indices- in Branch- cond- (buildProbTree @a left target df indicesL)- (buildProbTree @a right target df indicesR)--{- | Fit a TAO decision tree and return one @Expr Double@ per class.-- Each @(c, e)@ pair in the result map means: evaluate @e@ on a 'DataFrame'- row to get the predicted probability of class @c@. You can insert these- as new columns with 'derive' or evaluate them with 'interpret'.-- Example:- @- let pes = fitProbTree \@T.Text cfg (Col \"species\") trainDf- -- pes M.! \"setosa\" :: Expr Double- df' = M.foldlWithKey' (\\d cls e -> D.derive (cls <> \"_prob\") e d) testDf pes- @--}-fitProbTree ::- forall a.- (Columnable a, Ord a) =>- TreeConfig ->- Expr a -> -- target column, e.g. @Col \"label\"@- DataFrame ->- M.Map a (Expr Double)-fitProbTree cfg (Col target) df =- let- conds =- nubBy eqExpr $- numericConditions cfg (exclude [target] df)- ++ generateConditionsOld cfg (exclude [target] df)- initialTree = buildGreedyTree @a cfg (maxTreeDepth cfg) target conds df- indices = V.enumFromN 0 (nRows df)- optimizedTree = taoOptimize @a cfg target conds df indices initialTree- pruned = pruneDead optimizedTree- in- probExprs (buildProbTree @a pruned target df indices)-fitProbTree _ expr _ =- error $ "Cannot create prob tree for compound expression: " ++ show expr--{- | Convert a 'ProbTree' into one 'Expr Double' per class.-- Each @(c, e)@ pair means: evaluate @e@ on a 'DataFrame' row to get the- predicted probability of class @c@. You can insert these as new columns- with 'derive' or evaluate them with 'interpret'.-- Example:- @- let pt = fitProbTree \@T.Text cfg (Col \"species\") trainDf- pes = probExprs pt- -- pes M.! \"setosa\" :: Expr Double- df' = M.foldlWithKey' (\\d cls e -> D.derive (cls <> \"_prob\") e d) testDf pes- @--}-probExprs ::- forall a.- (Columnable a, Ord a) =>- ProbTree a ->- M.Map a (Expr Double)-probExprs tree =- let classes = nub (allClasses tree)- in M.fromList [(c, classExpr c tree) | c <- classes]- where- allClasses :: ProbTree a -> [a]- allClasses (Leaf m) = M.keys m- allClasses (Branch _ l r) = allClasses l ++ allClasses r-- classExpr :: a -> ProbTree a -> Expr Double- classExpr c (Leaf m) = Lit (M.findWithDefault 0.0 c m)- classExpr c (Branch cond l r) =- F.ifThenElse cond (classExpr c l) (classExpr c r)
− src/DataFrame/Display.hs
@@ -1,30 +0,0 @@-{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-}--{-# HLINT ignore "Use newtype instead of data" #-}-module DataFrame.Display where--import qualified Data.Text.IO as T-import DataFrame.Display.Terminal.PrettyPrint (RenderFormat (..))-import qualified DataFrame.Internal.DataFrame as D--data DisplayOptions = DisplayOptions- { {-- | Controls truncation on all three axes (rows, columns, cell width).- --- -- * 'Nothing' renders every row and column at full width.- -- * 'Just cfg' caps each axis whose limit in @cfg@ is positive.- --}- 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- { displayTruncate = Just D.defaultTruncateConfig{D.maxRows = 10}- }---- | Render a 'DataFrame' to stdout according to 'DisplayOptions'.-display :: DisplayOptions -> D.DataFrame -> IO ()-display opts = T.putStrLn . D.asTextWith Plain (displayTruncate opts)
− src/DataFrame/Display/Terminal/Colours.hs
@@ -1,14 +0,0 @@-module DataFrame.Display.Terminal.Colours where---- terminal color functions-red :: String -> String-red s = "\ESC[31m" ++ s ++ "\ESC[0m"--green :: String -> String-green s = "\ESC[32m" ++ s ++ "\ESC[0m"--brightGreen :: String -> String-brightGreen s = "\ESC[92m" ++ s ++ "\ESC[0m"--brightBlue :: String -> String-brightBlue s = "\ESC[94m" ++ s ++ "\ESC[0m"
− src/DataFrame/Display/Terminal/Plot.hs
@@ -1,648 +0,0 @@-{-# LANGUAGE AllowAmbiguousTypes #-}-{-# LANGUAGE ExplicitNamespaces #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE RankNTypes #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}--module DataFrame.Display.Terminal.Plot where--import Control.Monad-import qualified Data.Bifunctor-import qualified Data.List as L-import qualified Data.Map as M-import qualified Data.Text as T-import qualified Data.Text.IO as T-import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl))-import qualified Data.Vector as V-import qualified Data.Vector.Generic as VG-import qualified Data.Vector.Unboxed as VU-import Data.Word (Word8)-import DataFrame.Internal.Types-import GHC.Stack (HasCallStack)-import Type.Reflection (TypeRep, typeRep)--import DataFrame.Internal.Column (Column (..), Columnable, isNumeric)-import qualified DataFrame.Internal.Column as D-import DataFrame.Internal.DataFrame (DataFrame (..), getColumn)-import DataFrame.Internal.Expression-import DataFrame.Operations.Core-import qualified DataFrame.Operations.Subset as D-import Granite--data PlotConfig = PlotConfig- { plotType :: PlotType- , plotSettings :: Plot- }--data PlotType- = Histogram- | Scatter- | Line- | Bar- | BoxPlot- | Pie- | StackedBar- | Heatmap- deriving (Eq, Show)--defaultPlotConfig :: PlotType -> PlotConfig-defaultPlotConfig ptype =- PlotConfig- { plotType = ptype- , plotSettings = defPlot- }--plotHistogram :: (HasCallStack) => T.Text -> DataFrame -> IO ()-plotHistogram colName = plotHistogramWith colName 30 (defaultPlotConfig Histogram)--plotHistogramWith ::- (HasCallStack) => T.Text -> Int -> PlotConfig -> DataFrame -> IO ()-plotHistogramWith colName numBins config df = do- let values = extractNumericColumn colName df- (minVal, maxVal) = if null values then (0, 1) else (minimum values, maximum values)- T.putStrLn $ histogram (bins numBins minVal maxVal) values (plotSettings config)--plotScatter :: (HasCallStack) => T.Text -> T.Text -> DataFrame -> IO ()-plotScatter xCol yCol = plotScatterWith xCol yCol (defaultPlotConfig Scatter)--plotScatterWith ::- (HasCallStack) => T.Text -> T.Text -> PlotConfig -> DataFrame -> IO ()-plotScatterWith xCol yCol config df = do- let xVals = extractNumericColumn xCol df- yVals = extractNumericColumn yCol df- points = zip xVals yVals- T.putStrLn $ scatter [(xCol <> " vs " <> yCol, points)] (plotSettings config)--plotScatterBy ::- (HasCallStack) => T.Text -> T.Text -> T.Text -> DataFrame -> IO ()-plotScatterBy xCol yCol grouping = plotScatterByWith xCol yCol grouping (defaultPlotConfig Scatter)--plotScatterByWith ::- (HasCallStack) => T.Text -> T.Text -> T.Text -> PlotConfig -> DataFrame -> IO ()-plotScatterByWith xCol yCol grouping config df = do- let vals = extractStringColumn grouping df- let df' = insertColumn grouping (D.fromList vals) df- xs <- forM (L.nub vals) $ \col -> do- let filtered = D.filter (Col grouping) (== col) df'- xVals = extractNumericColumn xCol filtered- yVals = extractNumericColumn yCol filtered- points = zip xVals yVals- pure (col, points)- T.putStrLn $ scatter xs (plotSettings config)--plotLines :: (HasCallStack) => T.Text -> [T.Text] -> DataFrame -> IO ()-plotLines xAxis colNames = plotLinesWith xAxis colNames (defaultPlotConfig Line)--plotLinesWith ::- (HasCallStack) => T.Text -> [T.Text] -> PlotConfig -> DataFrame -> IO ()-plotLinesWith xAxis colNames config df = do- seriesData <- forM colNames $ \col -> do- let values = extractNumericColumn col df- indices = extractNumericColumn xAxis df- return (col, zip indices values)- T.putStrLn $ lineGraph seriesData (plotSettings config)--plotBoxPlots :: (HasCallStack) => [T.Text] -> DataFrame -> IO ()-plotBoxPlots colNames = plotBoxPlotsWith colNames (defaultPlotConfig BoxPlot)--plotBoxPlotsWith ::- (HasCallStack) => [T.Text] -> PlotConfig -> DataFrame -> IO ()-plotBoxPlotsWith colNames config df = do- boxData <- forM colNames $ \col -> do- let values = extractNumericColumn col df- return (col, values)- T.putStrLn $ boxPlot boxData (plotSettings config)--plotStackedBars :: (HasCallStack) => T.Text -> [T.Text] -> DataFrame -> IO ()-plotStackedBars categoryCol valueColumns = plotStackedBarsWith categoryCol valueColumns (defaultPlotConfig StackedBar)--plotStackedBarsWith ::- (HasCallStack) => T.Text -> [T.Text] -> PlotConfig -> DataFrame -> IO ()-plotStackedBarsWith categoryCol valueColumns config df = do- let categories = extractStringColumn categoryCol df- uniqueCategories = L.nub categories-- stackData <- forM uniqueCategories $ \cat -> do- let indices = [i | (i, c) <- zip [0 ..] categories, c == cat]- seriesData <- forM valueColumns $ \col -> do- let allValues = extractNumericColumn col df- values = [allValues !! i | i <- indices, i < length allValues]- return (col, sum values)- return (cat, seriesData)-- T.putStrLn $ stackedBars stackData (plotSettings config)--plotHeatmap :: (HasCallStack) => DataFrame -> IO ()-plotHeatmap = plotHeatmapWith (defaultPlotConfig Heatmap)--plotHeatmapWith :: (HasCallStack) => PlotConfig -> DataFrame -> IO ()-plotHeatmapWith config df = do- let numericCols = filter (isNumericColumn df) (columnNames df)- matrix = map (`extractNumericColumn` df) numericCols- T.putStrLn $ heatmap matrix (plotSettings config)--isNumericColumn :: DataFrame -> T.Text -> Bool-isNumericColumn df colName = maybe False isNumeric (getColumn colName df)--plotAllHistograms :: (HasCallStack) => DataFrame -> IO ()-plotAllHistograms df = do- let numericCols = filter (isNumericColumn df) (columnNames df)- forM_ numericCols $ \col -> do- T.putStrLn col- plotHistogram col df--plotCorrelationMatrix :: (HasCallStack) => DataFrame -> IO ()-plotCorrelationMatrix df = do- let numericCols = filter (isNumericColumn df) (columnNames df)- let correlations =- map- ( \col1 ->- map- ( \col2 ->- let- vals1 = extractNumericColumn col1 df- vals2 = extractNumericColumn col2 df- in- correlation vals1 vals2- )- numericCols- )- numericCols- print (zip [(0 :: Int) ..] numericCols)- T.putStrLn $ heatmap correlations (defPlot{plotTitle = "Correlation Matrix"})- where- correlation xs ys =- let n = fromIntegral $ length xs- meanX = sum xs / n- meanY = sum ys / n- covXY = sum [(x - meanX) * (y - meanY) | (x, y) <- zip xs ys] / n- stdX = sqrt $ sum [(x - meanX) ^ (2 :: Int) | x <- xs] / n- stdY = sqrt $ sum [(y - meanY) ^ (2 :: Int) | y <- ys] / n- in covXY / (stdX * stdY)--plotBars :: (HasCallStack) => T.Text -> DataFrame -> IO ()-plotBars colName = plotBarsWith colName Nothing (defaultPlotConfig Bar)--plotBarsWith ::- (HasCallStack) => T.Text -> Maybe T.Text -> PlotConfig -> DataFrame -> IO ()-plotBarsWith colName groupByCol config df =- case groupByCol of- Nothing -> plotSingleBars colName config df- Just grpCol -> plotGroupedBarsWith grpCol colName config df--plotSingleBars :: (HasCallStack) => T.Text -> PlotConfig -> DataFrame -> IO ()-plotSingleBars colName config df = do- let barData = getCategoricalCounts colName df- case barData of- Just counts -> do- let grouped = groupWithOther 10 counts- T.putStrLn $ bars grouped (plotSettings config)- Nothing -> do- let values = extractNumericColumn colName df- if length values > 20- then do- let labels = map (\i -> "Item " <> T.pack (show i)) [1 .. length values]- paired = zip labels values- grouped = groupWithOther 10 paired- T.putStrLn $ bars grouped (plotSettings config)- else do- let labels = map (\i -> "Item " <> T.pack (show i)) [1 .. length values]- T.putStrLn $ bars (zip labels values) (plotSettings config)--plotBarsTopN :: (HasCallStack) => Int -> T.Text -> DataFrame -> IO ()-plotBarsTopN n colName = plotBarsTopNWith n colName (defaultPlotConfig Bar)--plotBarsTopNWith ::- (HasCallStack) => Int -> T.Text -> PlotConfig -> DataFrame -> IO ()-plotBarsTopNWith n colName config df = do- let barData = getCategoricalCounts colName df- case barData of- Just counts -> do- let grouped = groupWithOther n counts- T.putStrLn $ bars grouped (plotSettings config)- Nothing -> do- let values = extractNumericColumn colName df- labels = map (\i -> "Item " <> T.pack (show i)) [1 .. length values]- paired = zip labels values- grouped = groupWithOther n paired- T.putStrLn $ bars grouped (plotSettings config)--plotGroupedBarsWith ::- (HasCallStack) => T.Text -> T.Text -> PlotConfig -> DataFrame -> IO ()-plotGroupedBarsWith = plotGroupedBarsWithN 10--plotGroupedBarsWithN ::- (HasCallStack) => Int -> T.Text -> T.Text -> PlotConfig -> DataFrame -> IO ()-plotGroupedBarsWithN n groupCol valCol config df = do- let colIsNumeric = isNumericColumnCheck valCol df-- if colIsNumeric- then do- let groups = extractStringColumn groupCol df- values = extractNumericColumn valCol df- m = M.fromListWith (+) (zip groups values)- grouped = map (\v -> (v, m M.! v)) groups- T.putStrLn $ bars grouped (plotSettings config)- else do- let groups = extractStringColumn groupCol df- vals = extractStringColumn valCol df- pairs = zip groups vals- counts =- M.toList $- M.fromListWith- (+)- [(g <> " - " <> v, 1 :: Int) | (g, v) <- pairs]- finalCounts = groupWithOther n [(k, fromIntegral v) | (k, v) <- counts]- T.putStrLn $ bars finalCounts (plotSettings config)--plotValueCounts :: (HasCallStack) => T.Text -> DataFrame -> IO ()-plotValueCounts colName = plotValueCountsWith colName 10 (defaultPlotConfig Bar)--plotValueCountsWith ::- (HasCallStack) => T.Text -> Int -> PlotConfig -> DataFrame -> IO ()-plotValueCountsWith colName maxBars config df = do- let counts = getCategoricalCounts colName df- case counts of- Just c -> do- let grouped = groupWithOther maxBars c- config' =- config- { plotSettings =- (plotSettings config)- { plotTitle =- if T.null (plotTitle (plotSettings config))- then "Value counts for " <> colName- else plotTitle (plotSettings config)- }- }- T.putStrLn $ bars grouped (plotSettings config')- Nothing -> error $ "Could not get value counts for column " ++ T.unpack colName--plotBarsWithPercentages :: (HasCallStack) => T.Text -> DataFrame -> IO ()-plotBarsWithPercentages colName df = do- let counts = getCategoricalCounts colName df- case counts of- Just c -> do- let total = sum (map snd c)- percentages =- [ (label <> " (" <> T.pack (show (round (100 * val / total) :: Int)) <> "%)", val)- | (label, val) <- c- ]- grouped = groupWithOther 10 percentages- T.putStrLn $ bars grouped (defPlot{plotTitle = "Distribution of " <> colName})- Nothing -> error $ "Could not get value counts for column " ++ T.unpack colName--smartPlotBars :: (HasCallStack) => T.Text -> DataFrame -> IO ()-smartPlotBars colName df = do- let counts = getCategoricalCounts colName df- case counts of- Just c -> do- let numUnique = length c- config =- (defaultPlotConfig Bar)- { plotSettings =- (plotSettings (defaultPlotConfig Bar))- { plotTitle = colName <> " (" <> T.pack (show numUnique) <> " unique values)"- }- }- if numUnique <= 12- then T.putStrLn $ bars c (plotSettings config)- else- if numUnique <= 20- then do- let grouped = groupWithOther 12 c- T.putStrLn $ bars grouped (plotSettings config)- else do- let grouped = groupWithOther 10 c- T.putStrLn $ bars grouped (plotSettings config)- Nothing -> plotBars colName df--plotCategoricalSummary :: (HasCallStack) => DataFrame -> IO ()-plotCategoricalSummary df = do- let cols = columnNames df- forM_ cols $ \col -> do- let counts = getCategoricalCounts col df- case counts of- Just c -> when (length c > 1) $ do- let numUnique = length c- putStrLn $- "\n=== " ++ T.unpack col ++ " (" ++ show numUnique ++ " unique values) ==="- if numUnique > 15 then plotBarsTopN 10 col df else plotBars col df- Nothing -> return ()--getCategoricalCounts ::- (HasCallStack) => T.Text -> DataFrame -> Maybe [(T.Text, Double)]-getCategoricalCounts colName df =- case M.lookup colName (columnIndices df) of- Nothing -> error $ "Column " ++ T.unpack colName ++ " not found"- Just idx ->- let col = columns df V.! idx- in case col of- BoxedColumn _ (vec :: V.Vector a) ->- Just (countBoxed (typeRep @a) vec)- UnboxedColumn _ (vec :: VU.Vector a) ->- Just (countUnboxed (typeRep @a) vec)- where- countBoxed ::- forall a. (Show a) => TypeRep a -> V.Vector a -> [(T.Text, Double)]- countBoxed tr vec- | Just Refl <- testEquality tr (typeRep @T.Text) = toPairs $ countValues vec- | Just Refl <- testEquality tr (typeRep @String) = toPairs $ countValues vec- | Just Refl <- testEquality tr (typeRep @Integer) = toPairs $ countValues vec- | Just Refl <- testEquality tr (typeRep @Int) = toPairs $ countValues vec- | Just Refl <- testEquality tr (typeRep @Double) = toPairs $ countValues vec- | Just Refl <- testEquality tr (typeRep @Float) = toPairs $ countValues vec- | Just Refl <- testEquality tr (typeRep @Bool) = toPairs $ countValues vec- | Just Refl <- testEquality tr (typeRep @Char) = toPairs $ countValues vec- | otherwise = countByShow $ V.toList vec-- countUnboxed ::- forall a. (Show a, VU.Unbox a) => TypeRep a -> VU.Vector a -> [(T.Text, Double)]- countUnboxed tr vec- | Just Refl <- testEquality tr (typeRep @Int) = toPairs $ countValuesUnboxed vec- | Just Refl <- testEquality tr (typeRep @Double) =- toPairs $ countValuesUnboxed vec- | Just Refl <- testEquality tr (typeRep @Float) =- toPairs $ countValuesUnboxed vec- | Just Refl <- testEquality tr (typeRep @Bool) =- toPairs $ countValuesUnboxed vec- | Just Refl <- testEquality tr (typeRep @Char) =- toPairs $ countValuesUnboxed vec- | Just Refl <- testEquality tr (typeRep @Word8) =- toPairs $ countValuesUnboxed vec- | otherwise = countByShow $ VU.toList vec-- toPairs :: (Show a) => [(a, Int)] -> [(T.Text, Double)]- toPairs = map (\(k, v) -> (T.pack (show k), fromIntegral v))-- countValues :: (Ord a) => V.Vector a -> [(a, Int)]- countValues vec = M.toList $ V.foldr' (\x acc -> M.insertWith (+) x 1 acc) M.empty vec-- countValuesUnboxed :: (Ord a, VU.Unbox a) => VU.Vector a -> [(a, Int)]- countValuesUnboxed vec = M.toList $ VU.foldr' (\x acc -> M.insertWith (+) x 1 acc) M.empty vec-- countByShow :: (Show a) => [a] -> [(T.Text, Double)]- countByShow xs =- map (Data.Bifunctor.bimap T.pack fromIntegral) $- M.toList $- L.foldl' (\acc x -> M.insertWith (+) (show x) (1 :: Int) acc) M.empty xs--isNumericColumnCheck :: T.Text -> DataFrame -> Bool-isNumericColumnCheck colName df = isNumericColumn df colName--extractStringColumn :: (HasCallStack) => T.Text -> DataFrame -> [T.Text]-extractStringColumn colName df =- case M.lookup colName (columnIndices df) of- Nothing -> error $ "Column " ++ T.unpack colName ++ " not found"- Just idx ->- let col = columns df V.! idx- in case col of- BoxedColumn _ vec -> V.toList $ V.map (T.pack . show) vec- UnboxedColumn _ vec -> V.toList $ VG.map (T.pack . show) (VG.convert vec)--extractNumericColumn :: (HasCallStack) => T.Text -> DataFrame -> [Double]-extractNumericColumn colName df =- case M.lookup colName (columnIndices df) of- Nothing -> error $ "Column " ++ T.unpack colName ++ " not found"- Just idx ->- let col = columns df V.! idx- in case col of- BoxedColumn _ vec -> vectorToDoubles vec- UnboxedColumn _ vec -> unboxedVectorToDoubles vec--vectorToDoubles :: forall a. (Columnable a, Show a) => V.Vector a -> [Double]-vectorToDoubles vec =- case testEquality (typeRep @a) (typeRep @Double) of- Just Refl -> V.toList vec- Nothing -> case sIntegral @a of- STrue -> V.toList $ V.map fromIntegral vec- SFalse -> case sFloating @a of- STrue -> V.toList $ V.map realToFrac vec- SFalse -> error $ "Column is not numeric (type: " ++ show (typeRep @a) ++ ")"--unboxedVectorToDoubles ::- forall a. (Columnable a, VU.Unbox a, Show a) => VU.Vector a -> [Double]-unboxedVectorToDoubles vec =- case testEquality (typeRep @a) (typeRep @Double) of- Just Refl -> VU.toList vec- Nothing -> case sIntegral @a of- STrue -> VU.toList $ VU.map fromIntegral vec- SFalse -> case sFloating @a of- STrue -> VU.toList $ VU.map realToFrac vec- SFalse -> error $ "Column is not numeric (type: " ++ show (typeRep @a) ++ ")"--groupWithOther :: Int -> [(T.Text, Double)] -> [(T.Text, Double)]-groupWithOther n items =- let sorted = L.sortOn (negate . snd) items- (topN, rest) = splitAt n sorted- otherSum = sum (map snd rest)- result =- if null rest || otherSum == 0- then topN- else topN ++ [("Other (" <> T.pack (show (length rest)) <> " items)", otherSum)]- in result--plotPie :: (HasCallStack) => T.Text -> Maybe T.Text -> DataFrame -> IO ()-plotPie valCol labelCol = plotPieWith valCol labelCol (defaultPlotConfig Pie)--plotPieWith ::- (HasCallStack) => T.Text -> Maybe T.Text -> PlotConfig -> DataFrame -> IO ()-plotPieWith valCol labelCol config df = do- let categoricalData = getCategoricalCounts valCol df- case categoricalData of- Just counts -> do- let grouped = groupWithOtherForPie 8 counts- T.putStrLn $ pie grouped (plotSettings config)- Nothing -> do- let values = extractNumericColumn valCol df- labels = case labelCol of- Nothing -> map (\i -> "Item " <> T.pack (show i)) [1 .. length values]- Just lCol -> extractStringColumn lCol df- let pieData = zip labels values- grouped =- if length pieData > 10- then groupWithOtherForPie 8 pieData- else pieData- T.putStrLn $ pie grouped (plotSettings config)--groupWithOtherForPie :: Int -> [(T.Text, Double)] -> [(T.Text, Double)]-groupWithOtherForPie n items =- let total = sum (map snd items)- sorted = L.sortOn (negate . snd) items- (topN, rest) = splitAt n sorted- otherSum = sum (map snd rest)- otherPct = round (100 * otherSum / total) :: Int- result =- if null rest || otherSum == 0- then topN- else- topN- ++ [- ( "Other ("- <> T.pack (show (length rest))- <> " items, "- <> T.pack (show otherPct)- <> "%)"- , otherSum- )- ]- in result--plotPieWithPercentages :: (HasCallStack) => T.Text -> DataFrame -> IO ()-plotPieWithPercentages colName = plotPieWithPercentagesConfig colName (defaultPlotConfig Pie)--plotPieWithPercentagesConfig ::- (HasCallStack) => T.Text -> PlotConfig -> DataFrame -> IO ()-plotPieWithPercentagesConfig colName config df = do- let counts = getCategoricalCounts colName df- case counts of- Just c -> do- let total = sum (map snd c)- withPct =- [ (label <> " (" <> T.pack (show (round (100 * val / total) :: Int)) <> "%)", val)- | (label, val) <- c- ]- grouped = groupWithOtherForPie 8 withPct- T.putStrLn $ pie grouped (plotSettings config)- Nothing -> do- let values = extractNumericColumn colName df- total = sum values- labels = map (\i -> "Item " <> T.pack (show i)) [1 .. length values]- withPct =- [ (label <> " (" <> T.pack (show (round (100 * val / total) :: Int)) <> "%)", val)- | (label, val) <- zip labels values- ]- grouped = groupWithOtherForPie 8 withPct- T.putStrLn $ pie grouped (plotSettings config)--plotPieTopN :: (HasCallStack) => Int -> T.Text -> DataFrame -> IO ()-plotPieTopN n colName = plotPieTopNWith n colName (defaultPlotConfig Pie)--plotPieTopNWith ::- (HasCallStack) => Int -> T.Text -> PlotConfig -> DataFrame -> IO ()-plotPieTopNWith n colName config df = do- let counts = getCategoricalCounts colName df- case counts of- Just c -> do- let grouped = groupWithOtherForPie n c- T.putStrLn $ pie grouped (plotSettings config)- Nothing -> do- let values = extractNumericColumn colName df- labels = map (\i -> "Item " <> T.pack (show i)) [1 .. length values]- paired = zip labels values- grouped = groupWithOtherForPie n paired- T.putStrLn $ pie grouped (plotSettings config)--smartPlotPie :: (HasCallStack) => T.Text -> DataFrame -> IO ()-smartPlotPie colName df = do- let counts = getCategoricalCounts colName df- case counts of- Just c -> do- let total = sum (map snd c)- significant = filter (\(_, v) -> v / total >= 0.01) c- config =- (defaultPlotConfig Pie)- { plotSettings =- (plotSettings (defaultPlotConfig Pie)){plotTitle = colName <> " Distribution"}- }- if length significant <= 6- then T.putStrLn $ pie significant (plotSettings config)- else- if length significant <= 10- then do- let grouped = groupWithOtherForPie 8 c- T.putStrLn $ pie grouped (plotSettings config)- else do- let grouped = groupWithOtherForPie 6 c- T.putStrLn $ pie grouped (plotSettings config)- Nothing -> plotPie colName Nothing df--plotPieGrouped :: (HasCallStack) => T.Text -> T.Text -> DataFrame -> IO ()-plotPieGrouped groupCol valCol = plotPieGroupedWith groupCol valCol (defaultPlotConfig Pie)--plotPieGroupedWith ::- (HasCallStack) => T.Text -> T.Text -> PlotConfig -> DataFrame -> IO ()-plotPieGroupedWith groupCol valCol config df = do- let colIsNumeric = isNumericColumnCheck valCol df-- if colIsNumeric- then do- let groups = extractStringColumn groupCol df- values = extractNumericColumn valCol df- grouped = M.toList $ M.fromListWith (+) (zip groups values)- finalGroups = groupWithOtherForPie 8 grouped- T.putStrLn $ pie finalGroups (plotSettings config)- else do- let groups = extractStringColumn groupCol df- vals = extractStringColumn valCol df- combined = zipWith (\g v -> g <> " - " <> v) groups vals- counts = M.toList $ M.fromListWith (+) [(c, 1 :: Int) | c <- combined]- finalCounts = groupWithOtherForPie 10 [(k, fromIntegral v) | (k, v) <- counts]- T.putStrLn $ pie finalCounts (plotSettings config)--plotPieComparison :: (HasCallStack) => [T.Text] -> DataFrame -> IO ()-plotPieComparison cols df = forM_ cols $ \col -> do- let counts = getCategoricalCounts col df- case counts of- Just c -> when (length c > 1 && length c <= 20) $ do- putStrLn $ "\n=== " ++ T.unpack col ++ " Distribution ==="- smartPlotPie col df- Nothing -> return ()--plotBinaryPie :: (HasCallStack) => T.Text -> DataFrame -> IO ()-plotBinaryPie colName df = do- let counts = getCategoricalCounts colName df- case counts of- Just c ->- if length c == 2- then do- let total = sum (map snd c)- withPct =- [ (label <> " (" <> T.pack (show (round (100 * val / total) :: Int)) <> "%)", val)- | (label, val) <- c- ]- T.putStrLn $ pie withPct defPlot- else- error $- "Column "- ++ T.unpack colName- ++ " is not binary (has "- ++ show (length c)- ++ " unique values)"- Nothing -> error $ "Column " ++ T.unpack colName ++ " is not categorical"--plotMarketShare :: (HasCallStack) => T.Text -> DataFrame -> IO ()-plotMarketShare colName = plotMarketShareWith colName (defaultPlotConfig Pie)--plotMarketShareWith ::- (HasCallStack) => T.Text -> PlotConfig -> DataFrame -> IO ()-plotMarketShareWith colName config df = do- let counts = getCategoricalCounts colName df- case counts of- Just c -> do- let total = sum (map snd c)- sorted = L.sortOn (negate . snd) c- significantShares = takeWhile (\(_, v) -> v / total >= 0.02) sorted- otherSum = sum [v | (_, v) <- c, v `notElem` map snd significantShares]-- formatShare (label, val) =- let pct = round (100 * val / total) :: Int- in (label <> " (" <> T.pack (show pct) <> "%)", val)-- shares = map formatShare significantShares- finalShares =- if otherSum > 0 && otherSum / total >= 0.01- then shares <> [("Others (<2% each)", otherSum)]- else shares-- let config' =- config- -- { plotSettings = (plotSettings config) {- -- plotTitle = colName <> ": market share"- -- }- -- }- T.putStrLn $ pie finalShares (plotSettings config')- Nothing -> error $ "Column " ++ T.unpack colName ++ " is not categorical"
− src/DataFrame/Display/Terminal/PrettyPrint.hs
@@ -1,104 +0,0 @@-{-# LANGUAGE OverloadedStrings #-}--module DataFrame.Display.Terminal.PrettyPrint where--import qualified Data.Text as T-import qualified Data.Vector as V--{- | 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.---- Adapted from: https://stackoverflow.com/questions/5929377/format-list-output-in-haskell--- a type for fill functions-type Filler = Int -> T.Text -> T.Text---- a type for describing table columns-data ColDesc t = ColDesc- { colTitleFill :: Filler- , colTitle :: T.Text- , colValueFill :: Filler- }---- 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.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) <> s--fillCenter :: Char -> Int -> T.Text -> T.Text-fillCenter c n s =- T.replicate l (T.singleton c) <> s <> T.replicate r (T.singleton c)- where- x = n - T.length s- l = x `div` 2- r = x - l---- functions that fill with spaces-left :: Int -> T.Text -> T.Text-left = fillLeft ' '--right :: Int -> T.Text -> T.Text-right = fillRight ' '--center :: Int -> T.Text -> T.Text-center = fillCenter ' '--{- | 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 =- 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 isMarkdown- then- wrapMd (fillCells colTitleFill consolidatedHeader)- : wrapMd separator- : map wrapMd rowLines- else- border- : fillCells colTitleFill consolidatedHeader- : separator- : fillCells colTitleFill types- : separator- : rowLines- in T.unlines outputLines
− src/DataFrame/Display/Web/Plot.hs
@@ -1,1080 +0,0 @@-{-# LANGUAGE AllowAmbiguousTypes #-}-{-# LANGUAGE ExplicitNamespaces #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE RankNTypes #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}--module DataFrame.Display.Web.Plot where--import Control.Monad-import qualified Data.Bifunctor-import Data.Char-import qualified Data.List as L-import qualified Data.Map as M-import qualified Data.Text as T-import qualified Data.Text.IO as T-import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl))-import qualified Data.Vector as V-import qualified Data.Vector.Generic as VG-import qualified Data.Vector.Unboxed as VU-import Data.Word (Word8)-import GHC.Stack (HasCallStack)-import System.Random (newStdGen, randomRs)-import Type.Reflection (TypeRep, typeRep)--import DataFrame.Internal.Column (Column (..), Columnable, isNumeric)-import qualified DataFrame.Internal.Column as D-import DataFrame.Internal.DataFrame (DataFrame (..), getColumn)-import DataFrame.Internal.Expression-import DataFrame.Internal.Types-import DataFrame.Operations.Core-import qualified DataFrame.Operations.Subset as D-import Numeric (showFFloat)-import System.Directory-import System.Info-import System.Process (- StdStream (NoStream),- createProcess,- proc,- std_err,- std_in,- std_out,- waitForProcess,- )--newtype HtmlPlot = HtmlPlot T.Text deriving (Show)--data PlotConfig = PlotConfig- { plotType :: PlotType- , plotTitle :: T.Text- , plotWidth :: Int- , plotHeight :: Int- , plotFile :: Maybe FilePath- }--data PlotType- = Histogram- | Scatter- | Line- | Bar- | BoxPlot- | Pie- | StackedBar- | Heatmap- deriving (Eq, Show)--defaultPlotConfig :: PlotType -> PlotConfig-defaultPlotConfig ptype =- PlotConfig- { plotType = ptype- , plotTitle = ""- , plotWidth = 600- , plotHeight = 400- , plotFile = Nothing- }--generateChartId :: IO T.Text-generateChartId = do- gen <- newStdGen- let randomWords =- filter- (\c -> c `elem` ([49 .. 57] ++ [65 .. 90] ++ [97 .. 122]))- (take 64 (randomRs (49, 126) gen :: [Int]))- return $ "chart_" <> T.pack (map chr randomWords)--wrapInHTML :: T.Text -> T.Text -> Int -> Int -> T.Text-wrapInHTML chartId content width height =- T.concat- [ "<canvas id=\""- , chartId- , "\" style=\"width:100%;max-width:"- , T.pack (show width)- , "px;height:"- , T.pack (show height)- , "px\"></canvas>\n"- , "<script src=\"https://cdnjs.cloudflare.com/ajax/libs/Chart.js/2.9.4/Chart.min.js\"></script>\n"- , "<script>\n"- , content- , "\n</script>\n"- ]--plotHistogram :: (HasCallStack) => T.Text -> DataFrame -> IO HtmlPlot-plotHistogram colName = plotHistogramWith colName 30 (defaultPlotConfig Histogram)--plotHistogramWith ::- (HasCallStack) => T.Text -> Int -> PlotConfig -> DataFrame -> IO HtmlPlot-plotHistogramWith colName numBins config df = do- chartId <- generateChartId- let values = extractNumericColumn colName df- (minVal, maxVal) = if null values then (0, 1) else (minimum values, maximum values)- binWidth = (maxVal - minVal) / fromIntegral numBins- bins = [minVal + fromIntegral i * binWidth | i <- [0 .. numBins - 1]]- counts = calculateHistogram values bins binWidth- precision = max 0 $ ceiling (negate $ logBase 10 binWidth)-- labels =- T.intercalate- ","- [ "\"" <> T.pack (showFFloat (Just precision) b "") <> "\""- | b <- bins- ]- dataPoints = T.intercalate "," [T.pack (show c) | c <- counts]-- chartTitle =- if T.null (plotTitle config)- then "Histogram of " <> colName- else plotTitle config-- jsCode =- T.concat- [ "setTimeout(function() { new Chart(\""- , chartId- , "\", {\n"- , " type: \"bar\",\n"- , " data: {\n"- , " labels: ["- , labels- , "],\n"- , " datasets: [{\n"- , " label: \""- , colName- , "\",\n"- , " data: ["- , dataPoints- , "],\n"- , " backgroundColor: \"rgba(75, 192, 192, 0.6)\",\n"- , " borderColor: \"rgba(75, 192, 192, 1)\",\n"- , " borderWidth: 1\n"- , " }]\n"- , " },\n"- , " options: {\n"- , " title: { display: true, text: \""- , chartTitle- , "\" },\n"- , " scales: {\n"- , " yAxes: [{ ticks: { beginAtZero: true } }]\n"- , " }\n"- , " }\n"- , "})}, 100);"- ]-- return $- HtmlPlot $- wrapInHTML chartId jsCode (plotWidth config) (plotHeight config)--calculateHistogram :: [Double] -> [Double] -> Double -> [Int]-calculateHistogram values bins binWidth =- let countBin b = length [v | v <- values, v >= b && v < b + binWidth]- in map countBin bins--plotScatter :: (HasCallStack) => T.Text -> T.Text -> DataFrame -> IO HtmlPlot-plotScatter xCol yCol = plotScatterWith xCol yCol (defaultPlotConfig Scatter)--plotScatterWith ::- (HasCallStack) => T.Text -> T.Text -> PlotConfig -> DataFrame -> IO HtmlPlot-plotScatterWith xCol yCol config df = do- chartId <- generateChartId- let xVals = extractNumericColumn xCol df- yVals = extractNumericColumn yCol df- points = zip xVals yVals-- dataPoints =- T.intercalate- ","- [ "{x:" <> T.pack (show x) <> ", y:" <> T.pack (show y) <> "}" | (x, y) <- points- ]- chartTitle =- if T.null (plotTitle config) then xCol <> " vs " <> yCol else plotTitle config-- jsCode =- T.concat- [ "setTimeout(function() { new Chart(\""- , chartId- , "\", {\n"- , " type: \"scatter\",\n"- , " data: {\n"- , " datasets: [{\n"- , " label: \""- , chartTitle- , "\",\n"- , " data: ["- , dataPoints- , "],\n"- , " pointRadius: 4,\n"- , " pointBackgroundColor: \"rgb(75, 192, 192)\"\n"- , " }]\n"- , " },\n"- , " options: {\n"- , " title: { display: true, text: \""- , chartTitle- , "\" },\n"- , " scales: {\n"- , " xAxes: [{ scaleLabel: { display: true, labelString: \""- , xCol- , "\" } }],\n"- , " yAxes: [{ scaleLabel: { display: true, labelString: \""- , yCol- , "\" } }]\n"- , " }\n"- , " }\n"- , "})}, 100);"- ]-- return $- HtmlPlot $- wrapInHTML chartId jsCode (plotWidth config) (plotHeight config)--plotScatterBy ::- (HasCallStack) => T.Text -> T.Text -> T.Text -> DataFrame -> IO HtmlPlot-plotScatterBy xCol yCol grouping = plotScatterByWith xCol yCol grouping (defaultPlotConfig Scatter)--plotScatterByWith ::- (HasCallStack) =>- T.Text -> T.Text -> T.Text -> PlotConfig -> DataFrame -> IO HtmlPlot-plotScatterByWith xCol yCol grouping config df = do- chartId <- generateChartId- let vals = extractStringColumn grouping df- df' = insertColumn grouping (D.fromList vals) df- uniqueVals = L.nub vals-- colors =- cycle- [ "rgb(255, 99, 132)"- , "rgb(54, 162, 235)"- , "rgb(255, 206, 86)"- , "rgb(75, 192, 192)"- , "rgb(153, 102, 255)"- , "rgb(255, 159, 64)"- ]-- datasets <- forM (zip uniqueVals colors) $ \(val, color) -> do- let filtered = D.filter (Col grouping) (== val) df'- xVals = extractNumericColumn xCol filtered- yVals = extractNumericColumn yCol filtered- points = zip xVals yVals- dataPoints =- T.intercalate- ","- [ "{x:" <> T.pack (show x) <> ", y:" <> T.pack (show y) <> "}" | (x, y) <- points- ]- return $- T.concat- [ " {\n"- , " label: \""- , val- , "\",\n"- , " data: ["- , dataPoints- , "],\n"- , " pointRadius: 4,\n"- , " pointBackgroundColor: \""- , color- , "\"\n"- , " }"- ]-- let datasetsStr = T.intercalate ",\n" datasets- chartTitle =- if T.null (plotTitle config)- then xCol <> " vs " <> yCol <> " by " <> grouping- else plotTitle config-- jsCode =- T.concat- [ "setTimeout(function() { new Chart(\""- , chartId- , "\", {\n"- , " type: \"scatter\",\n"- , " data: {\n"- , " datasets: [\n"- , datasetsStr- , "\n ]\n"- , " },\n"- , " options: {\n"- , " title: { display: true, text: \""- , chartTitle- , "\" },\n"- , " scales: {\n"- , " xAxes: [{ scaleLabel: { display: true, labelString: \""- , xCol- , "\" } }],\n"- , " yAxes: [{ scaleLabel: { display: true, labelString: \""- , yCol- , "\" } }]\n"- , " }\n"- , " }\n"- , "})}, 100);"- ]-- return $- HtmlPlot $- wrapInHTML chartId jsCode (plotWidth config) (plotHeight config)--plotLines :: (HasCallStack) => T.Text -> [T.Text] -> DataFrame -> IO HtmlPlot-plotLines xAxis colNames = plotLinesWith xAxis colNames (defaultPlotConfig Line)--plotLinesWith ::- (HasCallStack) => T.Text -> [T.Text] -> PlotConfig -> DataFrame -> IO HtmlPlot-plotLinesWith xAxis colNames config df = do- chartId <- generateChartId- let xValues = extractNumericColumn xAxis df- labels = T.intercalate "," [T.pack (show x) | x <- xValues]-- colors =- cycle- [ "rgb(255, 99, 132)"- , "rgb(54, 162, 235)"- , "rgb(255, 206, 86)"- , "rgb(75, 192, 192)"- , "rgb(153, 102, 255)"- , "rgb(255, 159, 64)"- ]-- datasets <- forM (zip colNames colors) $ \(col, color) -> do- let values = extractNumericColumn col df- dataPoints = T.intercalate "," [T.pack (show v) | v <- values]- return $- T.concat- [ " {\n"- , " label: \""- , col- , "\",\n"- , " data: ["- , dataPoints- , "],\n"- , " fill: false,\n"- , " borderColor: \""- , color- , "\",\n"- , " tension: 0.1\n"- , " }"- ]-- let datasetsStr = T.intercalate ",\n" datasets- chartTitle = if T.null (plotTitle config) then "Line Chart" else plotTitle config-- jsCode =- T.concat- [ "setTimeout(function() { new Chart(\""- , chartId- , "\", {\n"- , " type: \"line\",\n"- , " data: {\n"- , " labels: ["- , labels- , "],\n"- , " datasets: [\n"- , datasetsStr- , "\n ]\n"- , " },\n"- , " options: {\n"- , " title: { display: true, text: \""- , chartTitle- , "\" },\n"- , " scales: {\n"- , " xAxes: [{ scaleLabel: { display: true, labelString: \""- , xAxis- , "\" } }]\n"- , " }\n"- , " }\n"- , "})}, 100);"- ]-- return $- HtmlPlot $- wrapInHTML chartId jsCode (plotWidth config) (plotHeight config)--plotBars :: (HasCallStack) => T.Text -> DataFrame -> IO HtmlPlot-plotBars colName = plotBarsWith colName Nothing (defaultPlotConfig Bar)--plotBarsWith ::- (HasCallStack) =>- T.Text -> Maybe T.Text -> PlotConfig -> DataFrame -> IO HtmlPlot-plotBarsWith colName groupByCol config df =- case groupByCol of- Nothing -> plotSingleBars colName config df- Just grpCol -> plotGroupedBarsWith grpCol colName config df--plotSingleBars ::- (HasCallStack) => T.Text -> PlotConfig -> DataFrame -> IO HtmlPlot-plotSingleBars colName config df = do- chartId <- generateChartId- let barData = getCategoricalCounts colName df- case barData of- Just counts -> do- let grouped = groupWithOther 10 counts- labels = T.intercalate "," ["\"" <> label <> "\"" | (label, _) <- grouped]- dataPoints = T.intercalate "," [T.pack (show val) | (_, val) <- grouped]- chartTitle = if T.null (plotTitle config) then colName else plotTitle config-- jsCode =- T.concat- [ "setTimeout(function() { new Chart(\""- , chartId- , "\", {\n"- , " type: \"bar\",\n"- , " data: {\n"- , " labels: ["- , labels- , "],\n"- , " datasets: [{\n"- , " label: \"Count\",\n"- , " data: ["- , dataPoints- , "],\n"- , " backgroundColor: \"rgba(54, 162, 235, 0.6)\",\n"- , " borderColor: \"rgba(54, 162, 235, 1)\",\n"- , " borderWidth: 1\n"- , " }]\n"- , " },\n"- , " options: {\n"- , " title: { display: true, text: \""- , chartTitle- , "\" },\n"- , " scales: {\n"- , " yAxes: [{ ticks: { beginAtZero: true } }]\n"- , " }\n"- , " }\n"- , "})}, 100);"- ]- return $- HtmlPlot $- wrapInHTML chartId jsCode (plotWidth config) (plotHeight config)- Nothing -> do- let values = extractNumericColumn colName df- labels' =- if length values > 20- then take 20 ["Item " <> T.pack (show i) | i <- [(1 :: Int) ..]]- else ["Item " <> T.pack (show i) | i <- [1 .. length values]]- vals = if length values > 20 then take 20 values else values- labels = T.intercalate "," ["\"" <> label <> "\"" | label <- labels']- dataPoints = T.intercalate "," [T.pack (show val) | val <- vals]- chartTitle = if T.null (plotTitle config) then colName else plotTitle config-- jsCode =- T.concat- [ "setTimeout(function() { new Chart(\""- , chartId- , "\", {\n"- , " type: \"bar\",\n"- , " data: {\n"- , " labels: ["- , labels- , "],\n"- , " datasets: [{\n"- , " label: \"Value\",\n"- , " data: ["- , dataPoints- , "],\n"- , " backgroundColor: \"rgba(54, 162, 235, 0.6)\",\n"- , " borderColor: \"rgba(54, 162, 235, 1)\",\n"- , " borderWidth: 1\n"- , " }]\n"- , " },\n"- , " options: {\n"- , " title: { display: true, text: \""- , chartTitle- , "\" },\n"- , " scales: {\n"- , " yAxes: [{ ticks: { beginAtZero: true } }]\n"- , " }\n"- , " }\n"- , "})}, 100);"- ]- return $- HtmlPlot $- wrapInHTML chartId jsCode (plotWidth config) (plotHeight config)--plotPie :: (HasCallStack) => T.Text -> Maybe T.Text -> DataFrame -> IO HtmlPlot-plotPie valCol labelCol = plotPieWith valCol labelCol (defaultPlotConfig Pie)--plotPieWith ::- (HasCallStack) =>- T.Text -> Maybe T.Text -> PlotConfig -> DataFrame -> IO HtmlPlot-plotPieWith valCol labelCol config df = do- chartId <- generateChartId- let categoricalData = getCategoricalCounts valCol df- case categoricalData of- Just counts -> do- let grouped = groupWithOtherForPie 8 counts- labels = T.intercalate "," ["\"" <> label <> "\"" | (label, _) <- grouped]- dataPoints = T.intercalate "," [T.pack (show val) | (_, val) <- grouped]- colors = T.intercalate "," ["\"" <> c <> "\"" | c <- take (length grouped) pieColors]- chartTitle = if T.null (plotTitle config) then valCol else plotTitle config-- jsCode =- T.concat- [ "setTimeout(function() { new Chart(\""- , chartId- , "\", {\n"- , " type: \"pie\",\n"- , " data: {\n"- , " labels: ["- , labels- , "],\n"- , " datasets: [{\n"- , " data: ["- , dataPoints- , "],\n"- , " backgroundColor: ["- , colors- , "]\n"- , " }]\n"- , " },\n"- , " options: {\n"- , " title: { display: true, text: \""- , chartTitle- , "\" }\n"- , " }\n"- , "})}, 100);"- ]- return $- HtmlPlot $- wrapInHTML chartId jsCode (plotWidth config) (plotHeight config)- Nothing -> do- let values = extractNumericColumn valCol df- labels' = case labelCol of- Nothing -> map (\i -> "Item " <> T.pack (show i)) [1 .. length values]- Just lCol -> extractStringColumn lCol df- pieData = zip labels' values- grouped =- if length pieData > 10- then groupWithOtherForPie 8 pieData- else pieData- labels = T.intercalate "," ["\"" <> label <> "\"" | (label, _) <- grouped]- dataPoints = T.intercalate "," [T.pack (show val) | (_, val) <- grouped]- colors = T.intercalate "," ["\"" <> c <> "\"" | c <- take (length grouped) pieColors]- chartTitle = if T.null (plotTitle config) then valCol else plotTitle config-- jsCode =- T.concat- [ "setTimeout(function() { new Chart(\""- , chartId- , "\", {\n"- , " type: \"pie\",\n"- , " data: {\n"- , " labels: ["- , labels- , "],\n"- , " datasets: [{\n"- , " data: ["- , dataPoints- , "],\n"- , " backgroundColor: ["- , colors- , "]\n"- , " }]\n"- , " },\n"- , " options: {\n"- , " title: { display: true, text: \""- , chartTitle- , "\" }\n"- , " }\n"- , "})}, 100);"- ]- return $- HtmlPlot $- wrapInHTML chartId jsCode (plotWidth config) (plotHeight config)--pieColors :: [T.Text]-pieColors =- [ "rgb(255, 99, 132)"- , "rgb(54, 162, 235)"- , "rgb(255, 206, 86)"- , "rgb(75, 192, 192)"- , "rgb(153, 102, 255)"- , "rgb(255, 159, 64)"- , "rgb(201, 203, 207)"- , "rgb(255, 99, 71)"- , "rgb(60, 179, 113)"- , "rgb(238, 130, 238)"- ]--plotStackedBars ::- (HasCallStack) => T.Text -> [T.Text] -> DataFrame -> IO HtmlPlot-plotStackedBars categoryCol valueColumns = plotStackedBarsWith categoryCol valueColumns (defaultPlotConfig StackedBar)--plotStackedBarsWith ::- (HasCallStack) => T.Text -> [T.Text] -> PlotConfig -> DataFrame -> IO HtmlPlot-plotStackedBarsWith categoryCol valueColumns config df = do- chartId <- generateChartId- let categories = extractStringColumn categoryCol df- uniqueCategories = L.nub categories-- colors =- cycle- [ "rgb(255, 99, 132)"- , "rgb(54, 162, 235)"- , "rgb(255, 206, 86)"- , "rgb(75, 192, 192)"- , "rgb(153, 102, 255)"- , "rgb(255, 159, 64)"- ]-- datasets <- forM (zip valueColumns colors) $ \(col, color) -> do- dataVals <- forM uniqueCategories $ \cat -> do- let indices = [i | (i, c) <- zip [0 ..] categories, c == cat]- allValues = extractNumericColumn col df- values = [allValues !! i | i <- indices, i < length allValues]- return $ sum values- let dataPoints = T.intercalate "," [T.pack (show v) | v <- dataVals]- return $- T.concat- [ " {\n"- , " label: \""- , col- , "\",\n"- , " data: ["- , dataPoints- , "],\n"- , " backgroundColor: \""- , color- , "\"\n"- , " }"- ]-- let datasetsStr = T.intercalate ",\n" datasets- labels = T.intercalate "," ["\"" <> cat <> "\"" | cat <- uniqueCategories]- chartTitle = if T.null (plotTitle config) then "Stacked Bar Chart" else plotTitle config-- jsCode =- T.concat- [ "setTimeout(function() { new Chart(\""- , chartId- , "\", {\n"- , " type: \"bar\",\n"- , " data: {\n"- , " labels: ["- , labels- , "],\n"- , " datasets: [\n"- , datasetsStr- , "\n ]\n"- , " },\n"- , " options: {\n"- , " title: { display: true, text: \""- , chartTitle- , "\" },\n"- , " scales: {\n"- , " xAxes: [{ stacked: true }],\n"- , " yAxes: [{ stacked: true, ticks: { beginAtZero: true } }]\n"- , " }\n"- , " }\n"- , "})}, 100);"- ]-- return $- HtmlPlot $- wrapInHTML chartId jsCode (plotWidth config) (plotHeight config)--plotBoxPlots :: (HasCallStack) => [T.Text] -> DataFrame -> IO HtmlPlot-plotBoxPlots colNames = plotBoxPlotsWith colNames (defaultPlotConfig BoxPlot)--plotBoxPlotsWith ::- (HasCallStack) => [T.Text] -> PlotConfig -> DataFrame -> IO HtmlPlot-plotBoxPlotsWith colNames config df = do- chartId <- generateChartId- boxData <- forM colNames $ \col -> do- let values = extractNumericColumn col df- sorted = L.sort values- n = length values- q1 = sorted !! (n `div` 4)- median = sorted !! (n `div` 2)- q3 = sorted !! (3 * n `div` 4)- minVal = minimum values- maxVal = maximum values- return (col, minVal, q1, median, q3, maxVal)-- let labels = T.intercalate "," ["\"" <> col <> "\"" | (col, _, _, _, _, _) <- boxData]- medians = T.intercalate "," [T.pack (show med) | (_, _, _, med, _, _) <- boxData]- chartTitle = if T.null (plotTitle config) then "Box Plot" else plotTitle config-- jsCode =- T.concat- [ "setTimeout(function() { new Chart(\""- , chartId- , "\", {\n"- , " type: \"bar\",\n"- , " data: {\n"- , " labels: ["- , labels- , "],\n"- , " datasets: [{\n"- , " label: \"Median\",\n"- , " data: ["- , medians- , "],\n"- , " backgroundColor: \"rgba(75, 192, 192, 0.6)\",\n"- , " borderColor: \"rgba(75, 192, 192, 1)\",\n"- , " borderWidth: 1\n"- , " }]\n"- , " },\n"- , " options: {\n"- , " title: { display: true, text: \""- , chartTitle- , " (showing medians)\" },\n"- , " scales: {\n"- , " yAxes: [{ ticks: { beginAtZero: true } }]\n"- , " }\n"- , " }\n"- , "})}, 100);"- ]-- return $- HtmlPlot $- wrapInHTML chartId jsCode (plotWidth config) (plotHeight config)--plotGroupedBarsWith ::- (HasCallStack) => T.Text -> T.Text -> PlotConfig -> DataFrame -> IO HtmlPlot-plotGroupedBarsWith = plotGroupedBarsWithN 10--plotGroupedBarsWithN ::- (HasCallStack) =>- Int -> T.Text -> T.Text -> PlotConfig -> DataFrame -> IO HtmlPlot-plotGroupedBarsWithN n groupCol valCol config df = do- chartId <- generateChartId- let colIsNumeric = isNumericColumnCheck valCol df-- if colIsNumeric- then do- let groups = extractStringColumn groupCol df- values = extractNumericColumn valCol df- m = M.fromListWith (+) (zip groups values)- grouped = map (\v -> (v, m M.! v)) groups- labels = T.intercalate "," ["\"" <> label <> "\"" | (label, _) <- grouped]- dataPoints = T.intercalate "," [T.pack (show val) | (_, val) <- grouped]- chartTitle =- if T.null (plotTitle config)- then groupCol <> " by " <> valCol- else plotTitle config-- jsCode =- T.concat- [ "setTimeout(function() { new Chart(\""- , chartId- , "\", {\n"- , " type: \"bar\",\n"- , " data: {\n"- , " labels: ["- , labels- , "],\n"- , " datasets: [{\n"- , " label: \""- , valCol- , "\",\n"- , " data: ["- , dataPoints- , "],\n"- , " backgroundColor: \"rgba(54, 162, 235, 0.6)\",\n"- , " borderColor: \"rgba(54, 162, 235, 1)\",\n"- , " borderWidth: 1\n"- , " }]\n"- , " },\n"- , " options: {\n"- , " title: { display: true, text: \""- , chartTitle- , "\" },\n"- , " scales: {\n"- , " yAxes: [{ ticks: { beginAtZero: true } }]\n"- , " }\n"- , " }\n"- , "})}, 100);"- ]- return $- HtmlPlot $- wrapInHTML chartId jsCode (plotWidth config) (plotHeight config)- else do- let groups = extractStringColumn groupCol df- vals = extractStringColumn valCol df- pairs = zip groups vals- counts =- M.toList $- M.fromListWith- (+)- [(g <> " - " <> v, 1 :: Int) | (g, v) <- pairs]- finalCounts = groupWithOther n [(k, fromIntegral v) | (k, v) <- counts]- labels = T.intercalate "," ["\"" <> label <> "\"" | (label, _) <- finalCounts]- dataPoints = T.intercalate "," [T.pack (show val) | (_, val) <- finalCounts]- chartTitle =- if T.null (plotTitle config)- then groupCol <> " by " <> valCol- else plotTitle config-- jsCode =- T.concat- [ "setTimeout(function() { new Chart(\""- , chartId- , "\", {\n"- , " type: \"bar\",\n"- , " data: {\n"- , " labels: ["- , labels- , "],\n"- , " datasets: [{\n"- , " label: \"Count\",\n"- , " data: ["- , dataPoints- , "],\n"- , " backgroundColor: \"rgba(54, 162, 235, 0.6)\",\n"- , " borderColor: \"rgba(54, 162, 235, 1)\",\n"- , " borderWidth: 1\n"- , " }]\n"- , " },\n"- , " options: {\n"- , " title: { display: true, text: \""- , chartTitle- , "\" },\n"- , " scales: {\n"- , " yAxes: [{ ticks: { beginAtZero: true } }]\n"- , " }\n"- , " }\n"- , "})}, 100);"- ]- return $- HtmlPlot $- wrapInHTML chartId jsCode (plotWidth config) (plotHeight config)---- TODO: Move these helpers to a common module.--isNumericColumn :: DataFrame -> T.Text -> Bool-isNumericColumn df colName = maybe False isNumeric (getColumn colName df)--isNumericColumnCheck :: T.Text -> DataFrame -> Bool-isNumericColumnCheck colName df = isNumericColumn df colName--extractStringColumn :: (HasCallStack) => T.Text -> DataFrame -> [T.Text]-extractStringColumn colName df =- case M.lookup colName (columnIndices df) of- Nothing -> error $ "Column " ++ T.unpack colName ++ " not found"- Just idx ->- let col = columns df V.! idx- in case col of- BoxedColumn _ (vec :: V.Vector a) -> case testEquality (typeRep @a) (typeRep @T.Text) of- Just Refl -> V.toList vec- Nothing -> V.toList $ V.map (T.pack . show) vec- UnboxedColumn _ vec -> V.toList $ VG.map (T.pack . show) (VG.convert vec)--extractNumericColumn :: (HasCallStack) => T.Text -> DataFrame -> [Double]-extractNumericColumn colName df =- case M.lookup colName (columnIndices df) of- Nothing -> error $ "Column " ++ T.unpack colName ++ " not found"- Just idx ->- let col = columns df V.! idx- in case col of- BoxedColumn _ vec -> vectorToDoubles vec- UnboxedColumn _ vec -> unboxedVectorToDoubles vec--vectorToDoubles :: forall a. (Columnable a, Show a) => V.Vector a -> [Double]-vectorToDoubles vec =- case testEquality (typeRep @a) (typeRep @Double) of- Just Refl -> V.toList vec- Nothing -> case sIntegral @a of- STrue -> V.toList $ V.map fromIntegral vec- SFalse -> case sFloating @a of- STrue -> V.toList $ V.map realToFrac vec- SFalse -> error $ "Column is not numeric (type: " ++ show (typeRep @a) ++ ")"--unboxedVectorToDoubles ::- forall a. (Columnable a, VU.Unbox a, Show a) => VU.Vector a -> [Double]-unboxedVectorToDoubles vec =- case testEquality (typeRep @a) (typeRep @Double) of- Just Refl -> VU.toList vec- Nothing -> case sIntegral @a of- STrue -> VU.toList $ VU.map fromIntegral vec- SFalse -> case sFloating @a of- STrue -> VU.toList $ VU.map realToFrac vec- SFalse -> error $ "Column is not numeric (type: " ++ show (typeRep @a) ++ ")"--getCategoricalCounts ::- (HasCallStack) => T.Text -> DataFrame -> Maybe [(T.Text, Double)]-getCategoricalCounts colName df =- case M.lookup colName (columnIndices df) of- Nothing -> error $ "Column " ++ T.unpack colName ++ " not found"- Just idx ->- let col = columns df V.! idx- in case col of- BoxedColumn _ (vec :: V.Vector a) ->- Just (countBoxed (typeRep @a) vec)- UnboxedColumn _ (vec :: VU.Vector a) ->- Just (countUnboxed (typeRep @a) vec)- where- countBoxed ::- forall a. (Show a) => TypeRep a -> V.Vector a -> [(T.Text, Double)]- countBoxed tr vec- | Just Refl <- testEquality tr (typeRep @T.Text) = toPairsText $ countValues vec- | Just Refl <- testEquality tr (typeRep @String) = toPairs $ countValues vec- | Just Refl <- testEquality tr (typeRep @Integer) = toPairs $ countValues vec- | Just Refl <- testEquality tr (typeRep @Int) = toPairs $ countValues vec- | Just Refl <- testEquality tr (typeRep @Double) = toPairs $ countValues vec- | Just Refl <- testEquality tr (typeRep @Float) = toPairs $ countValues vec- | Just Refl <- testEquality tr (typeRep @Bool) = toPairs $ countValues vec- | Just Refl <- testEquality tr (typeRep @Char) = toPairs $ countValues vec- | otherwise = countByShow $ V.toList vec-- countUnboxed ::- forall a. (Show a, VU.Unbox a) => TypeRep a -> VU.Vector a -> [(T.Text, Double)]- countUnboxed tr vec- | Just Refl <- testEquality tr (typeRep @Int) = toPairs $ countValuesUnboxed vec- | Just Refl <- testEquality tr (typeRep @Double) =- toPairs $ countValuesUnboxed vec- | Just Refl <- testEquality tr (typeRep @Float) =- toPairs $ countValuesUnboxed vec- | Just Refl <- testEquality tr (typeRep @Bool) =- toPairs $ countValuesUnboxed vec- | Just Refl <- testEquality tr (typeRep @Char) =- toPairs $ countValuesUnboxed vec- | Just Refl <- testEquality tr (typeRep @Word8) =- toPairs $ countValuesUnboxed vec- | otherwise = countByShow $ VU.toList vec-- toPairs :: (Show a) => [(a, Int)] -> [(T.Text, Double)]- toPairs = map (\(k, v) -> (T.pack (show k), fromIntegral v))-- toPairsText :: [(T.Text, Int)] -> [(T.Text, Double)]- toPairsText = map (Data.Bifunctor.second fromIntegral)-- countValues :: (Ord a) => V.Vector a -> [(a, Int)]- countValues vec = M.toList $ V.foldr' (\x acc -> M.insertWith (+) x 1 acc) M.empty vec-- countValuesUnboxed :: (Ord a, VU.Unbox a) => VU.Vector a -> [(a, Int)]- countValuesUnboxed vec = M.toList $ VU.foldr' (\x acc -> M.insertWith (+) x 1 acc) M.empty vec-- countByShow :: (Show a) => [a] -> [(T.Text, Double)]- countByShow xs =- map (Data.Bifunctor.bimap T.pack fromIntegral) $- M.toList $- L.foldl' (\acc x -> M.insertWith (+) (show x) (1 :: Int) acc) M.empty xs--groupWithOther :: Int -> [(T.Text, Double)] -> [(T.Text, Double)]-groupWithOther n items =- let sorted = L.sortOn (negate . snd) items- (topN, rest) = splitAt n sorted- otherSum = sum (map snd rest)- result =- if null rest || otherSum == 0- then topN- else topN ++ [("Other (" <> T.pack (show (length rest)) <> " items)", otherSum)]- in result--groupWithOtherForPie :: Int -> [(T.Text, Double)] -> [(T.Text, Double)]-groupWithOtherForPie n items =- let total = sum (map snd items)- sorted = L.sortOn (negate . snd) items- (topN, rest) = splitAt n sorted- otherSum = sum (map snd rest)- otherPct = round (100 * otherSum / total) :: Int- result =- if null rest || otherSum == 0- then topN- else- topN- ++ [- ( "Other ("- <> T.pack (show (length rest))- <> " items, "- <> T.pack (show otherPct)- <> "%)"- , otherSum- )- ]- in result--plotBarsTopN :: (HasCallStack) => Int -> T.Text -> DataFrame -> IO HtmlPlot-plotBarsTopN n colName = plotBarsTopNWith n colName (defaultPlotConfig Bar)--plotBarsTopNWith ::- (HasCallStack) => Int -> T.Text -> PlotConfig -> DataFrame -> IO HtmlPlot-plotBarsTopNWith n colName config df = do- let config' = config{plotTitle = plotTitle config <> " (Top " <> T.pack (show n) <> ")"}- plotBarsWith colName Nothing config' df--plotValueCounts :: (HasCallStack) => T.Text -> DataFrame -> IO HtmlPlot-plotValueCounts colName = plotValueCountsWith colName 10 (defaultPlotConfig Bar)--plotValueCountsWith ::- (HasCallStack) => T.Text -> Int -> PlotConfig -> DataFrame -> IO HtmlPlot-plotValueCountsWith colName maxBars config df = do- let config' = config{plotTitle = "Value counts for " <> colName}- plotBarsTopNWith maxBars colName config' df--plotAllHistograms :: (HasCallStack) => DataFrame -> IO HtmlPlot-plotAllHistograms df = do- let numericCols = filter (isNumericColumn df) (columnNames df)- xs <- forM numericCols $ \col -> do- plotHistogram col df- let allPlots = L.foldl' (\acc (HtmlPlot contents) -> acc <> "\n" <> contents) "" xs- return (HtmlPlot allPlots)--plotCategoricalSummary :: (HasCallStack) => DataFrame -> IO HtmlPlot-plotCategoricalSummary df = do- let cols = columnNames df- xs <- forM cols $ \col -> do- let counts = getCategoricalCounts col df- case counts of- Just c -> do- if length c > 1- then- ( do- let numUnique = length c- putStrLn $- "\n<!-- " ++ T.unpack col ++ " (" ++ show numUnique ++ " unique values) -->"- if numUnique > 15 then plotBarsTopN 10 col df else plotBars col df- )- else return (HtmlPlot "")- Nothing -> return (HtmlPlot "")- let allPlots = L.foldl' (\acc (HtmlPlot contents) -> acc <> "\n" <> contents) "" xs- return (HtmlPlot allPlots)--plotBarsWithPercentages :: (HasCallStack) => T.Text -> DataFrame -> IO HtmlPlot-plotBarsWithPercentages colName df = do- let config = (defaultPlotConfig Bar){plotTitle = "Distribution of " <> colName}- plotBarsWith colName Nothing config df--smartPlotBars :: (HasCallStack) => T.Text -> DataFrame -> IO HtmlPlot-smartPlotBars colName df = do- let counts = getCategoricalCounts colName df- case counts of- Just c -> do- let numUnique = length c- config =- (defaultPlotConfig Bar)- { plotTitle = colName <> " (" <> T.pack (show numUnique) <> " unique values)"- }- if numUnique <= 12- then plotBarsWith colName Nothing config df- else plotBarsTopNWith 10 colName config df- Nothing -> plotBars colName df--showInDefaultBrowser :: HtmlPlot -> IO ()-showInDefaultBrowser (HtmlPlot p) = do- plotId <- generateChartId- home <- getHomeDirectory- let operatingSystem = os- let path = "plot-" <> T.unpack plotId <> ".html"-- let fullPath =- if operatingSystem == "mingw32"- then home <> "\\" <> path- else home <> "/" <> path- putStr "Saving plot to: "- putStrLn fullPath- T.writeFile fullPath p- case operatingSystem of- "mingw32" -> openFileSilently "start" fullPath- "darwin" -> openFileSilently "open" fullPath- _ -> openFileSilently "xdg-open" fullPath- pure ()--openFileSilently :: FilePath -> FilePath -> IO ()-openFileSilently program path = do- (_, _, _, ph) <-- createProcess- (proc program [path])- { std_in = NoStream- , std_out = NoStream- , std_err = NoStream- }- void (waitForProcess ph)
− src/DataFrame/Errors.hs
@@ -1,188 +0,0 @@-{-# LANGUAGE DeriveAnyClass #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE InstanceSigs #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE RankNTypes #-}--module DataFrame.Errors where--import qualified Data.Text as T-import qualified Data.Vector.Unboxed as VU--import Control.Exception-import Data.Array-import qualified Data.List as L-import Data.Typeable (Typeable)-import DataFrame.Display.Terminal.Colours-import Type.Reflection (TypeRep)--data TypeErrorContext a b = MkTypeErrorContext- { userType :: Either String (TypeRep a)- , expectedType :: Either String (TypeRep b)- , errorColumnName :: Maybe String- , callingFunctionName :: Maybe String- }--data DataFrameException where- TypeMismatchException ::- forall a b.- (Typeable a, Typeable b) =>- TypeErrorContext a b ->- DataFrameException- AggregatedAndNonAggregatedException :: T.Text -> T.Text -> DataFrameException- ColumnsNotFoundException :: [T.Text] -> T.Text -> [T.Text] -> DataFrameException- EmptyDataSetException :: T.Text -> DataFrameException- InternalException :: T.Text -> DataFrameException- NonColumnReferenceException :: T.Text -> DataFrameException- UnaggregatedException :: T.Text -> DataFrameException- WrongQuantileNumberException :: Int -> DataFrameException- WrongQuantileIndexException :: VU.Vector Int -> Int -> DataFrameException- deriving (Exception)--instance Show DataFrameException where- show :: DataFrameException -> String- show (TypeMismatchException context) =- let- errorString =- typeMismatchError- (either id show (userType context))- (either id show (expectedType context))- in- addCallPointInfo- (errorColumnName context)- (callingFunctionName context)- errorString- show (ColumnsNotFoundException columnNames callPoint availableColumns) = columnsNotFound columnNames callPoint availableColumns- show (EmptyDataSetException callPoint) = emptyDataSetError callPoint- show (WrongQuantileNumberException q) = wrongQuantileNumberError q- show (WrongQuantileIndexException qs q) = wrongQuantileIndexError qs q- show (InternalException msg) = "Internal error: " ++ T.unpack msg- show (NonColumnReferenceException msg) = "Expression must be a column reference in: " ++ T.unpack msg- show (UnaggregatedException expr) = "Expression is not fully aggregated: " ++ T.unpack expr- show (AggregatedAndNonAggregatedException expr1 expr2) =- "Cannot combine aggregated and non-aggregated expressions: \n"- ++ T.unpack expr1- ++ "\n"- ++ T.unpack expr2--columnNotFound :: T.Text -> T.Text -> [T.Text] -> String-columnNotFound missingColumn = columnsNotFound [missingColumn]--columnsNotFound :: [T.Text] -> T.Text -> [T.Text] -> String-columnsNotFound missingColumns callPoint availableColumns =- red "\n\n[ERROR] "- ++ missingColumnsLabel missingColumns- ++ ": "- ++ T.unpack (T.intercalate ", " missingColumns)- ++ " for operation "- ++ T.unpack callPoint- ++ formatSuggestions missingColumns availableColumns- ++ "\n\n"- where- missingColumnsLabel [_] = "Column not found"- missingColumnsLabel _ = "Columns not found"-- formatSuggestions [missingColumn] columns =- case guessColumnName missingColumn columns of- "" -> ""- guessed ->- "\n\tDid you mean "- ++ T.unpack guessed- ++ "?"- formatSuggestions names columns =- case traverse (`suggestColumnName` columns) names of- Just guessedColumns- | not (null guessedColumns) ->- "\n\tDid you mean "- ++ formatColumnSuggestions guessedColumns- ++ "?"- _ -> ""-- suggestColumnName missingColumn columns = case guessColumnName missingColumn columns of- "" -> Nothing- guessed -> Just guessed-- formatColumnSuggestions guessedColumns =- "["- ++ L.intercalate ", " (map (show . T.unpack) guessedColumns)- ++ "]"--typeMismatchError :: String -> String -> String-typeMismatchError givenType expType =- red $- red "\n\n[Error]: Type Mismatch"- ++ "\n\tWhile running your code I tried to "- ++ "get a column of type: "- ++ red (show givenType)- ++ " but the column in the dataframe was actually of type: "- ++ green (show expType)--emptyDataSetError :: T.Text -> String-emptyDataSetError callPoint =- red "\n\n[ERROR] "- ++ T.unpack callPoint- ++ " cannot be called on empty data sets"--wrongQuantileNumberError :: Int -> String-wrongQuantileNumberError q =- red "\n\n[ERROR] "- ++ "Quantile number q should satisfy "- ++ "q >= 2, but here q is "- ++ show q--wrongQuantileIndexError :: VU.Vector Int -> Int -> String-wrongQuantileIndexError qs q =- red "\n\n[ERROR] "- ++ "For quantile number q, "- ++ "each quantile index i "- ++ "should satisfy 0 <= i <= q, "- ++ "but here q is "- ++ show q- ++ " and indexes are "- ++ show qs--addCallPointInfo :: Maybe String -> Maybe String -> String -> String-addCallPointInfo (Just name) (Just cp) err =- err- ++ ( "\n\tThis happened when calling function "- ++ brightGreen cp- ++ " on "- ++ brightGreen name- )-addCallPointInfo Nothing (Just cp) err =- err- ++ ( "\n\tThis happened when calling function "- ++ brightGreen cp- )-addCallPointInfo (Just name) Nothing err =- err- ++ ( "\n\tOn "- ++ name- ++ "\n\n"- )-addCallPointInfo Nothing Nothing err = err--guessColumnName :: T.Text -> [T.Text] -> T.Text-guessColumnName userInput columns = case map (\k -> (editDistance userInput k, k)) columns of- [] -> ""- res -> (snd . minimum) res--editDistance :: T.Text -> T.Text -> Int-editDistance xs ys = table ! (m, n)- where- (m, n) = (T.length xs, T.length ys)- x = array (1, m) (zip [1 ..] (T.unpack xs))- y = array (1, n) (zip [1 ..] (T.unpack ys))-- table :: Array (Int, Int) Int- table = array bnds [(ij, dist ij) | ij <- range bnds]- bnds = ((0, 0), (m, n))-- dist (0, j) = j- dist (i, 0) = i- dist (i, j) =- minimum- [ table ! (i - 1, j) + 1- , table ! (i, j - 1) + 1- , if x ! i == y ! j then table ! (i - 1, j - 1) else 1 + table ! (i - 1, j - 1)- ]
− src/DataFrame/Functions.hs
@@ -1,676 +0,0 @@-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE FlexibleInstances #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE IncoherentInstances #-}-{-# LANGUAGE MultiParamTypeClasses #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE RankNTypes #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}-{-# LANGUAGE TypeOperators #-}-{-# LANGUAGE UndecidableInstances #-}--module DataFrame.Functions (module DataFrame.Functions, module DataFrame.Operators) where--import DataFrame.Internal.Column-import DataFrame.Internal.Expression-import DataFrame.Internal.Statistics--import Control.Applicative-import qualified Data.Char as Char-import Data.Either-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.Text as T-import Data.Time-import qualified Data.Vector as V-import qualified Data.Vector.Unboxed as VU--import DataFrame.Internal.Nullable (- BaseType,- NullLift1Op (applyNull1),- NullLift1Result,- NullLift2Op (applyNull2),- NullLift2Result,- )-import DataFrame.Operators-import Text.Regex.TDFA-import Prelude hiding (maximum, minimum)-import Prelude as P--lift :: (Columnable a, Columnable b) => (a -> b) -> Expr a -> Expr b-lift f =- Unary (MkUnaryOp{unaryFn = f, unaryName = "unaryUdf", unarySymbol = Nothing})--lift2 ::- (Columnable c, Columnable b, Columnable a) =>- (c -> b -> a) ->- Expr c ->- Expr b ->- Expr a-lift2 f =- Binary- ( MkBinaryOp- { binaryFn = f- , binaryName = "binaryUdf"- , binarySymbol = Nothing- , binaryCommutative = False- , binaryPrecedence = 0- }- )--{- | Lift a unary function over a nullable or non-nullable column expression.-When the input is @Maybe a@, 'Nothing' short-circuits (like 'fmap').-When the input is plain @a@, the function is applied directly.--The return type is inferred via 'NullLift1Result': no annotation needed.--}-nullLift ::- (NullLift1Op a r (NullLift1Result a r), Columnable (NullLift1Result a r)) =>- (BaseType a -> r) ->- Expr a ->- Expr (NullLift1Result a r)-nullLift f =- Unary- (MkUnaryOp{unaryFn = applyNull1 f, unaryName = "nullLift", unarySymbol = Nothing})--{- | Lift a binary function over nullable or non-nullable column expressions.-Any 'Nothing' operand short-circuits to 'Nothing' in the result.--The return type is inferred via 'NullLift2Result': no annotation needed.--}-nullLift2 ::- (NullLift2Op a b r (NullLift2Result a b r), Columnable (NullLift2Result a b r)) =>- (BaseType a -> BaseType b -> r) ->- Expr a ->- Expr b ->- Expr (NullLift2Result a b r)-nullLift2 f =- Binary- ( MkBinaryOp- { binaryFn = applyNull2 f- , binaryName = "nullLift2"- , binarySymbol = Nothing- , binaryCommutative = False- , binaryPrecedence = 0- }- )--{- | Lenient numeric \/ text coercion returning @Maybe a@. Looks up column-@name@ and coerces its values to @a@. Values that cannot be converted-(parse failures, type mismatches) become 'Nothing'; successfully converted-values are wrapped in 'Just'. Existing 'Nothing' in optional source columns-stays as 'Nothing'.--}-cast :: forall a. (Columnable a, Read a) => T.Text -> Expr (Maybe a)-cast colName = CastWith colName "cast" (either (const Nothing) Just)--{- | Lenient coercion that substitutes a default for unconvertible values.-Looks up column @name@, coerces its values to @a@, and uses @def@ wherever-conversion fails or the source value is 'Nothing'.--}-castWithDefault :: forall a. (Columnable a, Read a) => a -> T.Text -> Expr a-castWithDefault def colName =- CastWith colName ("castWithDefault:" <> T.pack (show def)) (fromRight def)--{- | Lenient coercion returning @Either T.Text a@. Successfully converted-values are 'Right'; values that cannot be parsed are kept as 'Left' with-their original string representation, so the caller can inspect or handle-them downstream. Existing 'Nothing' in optional source columns becomes-@Left \"null\"@.--}-castEither ::- forall a. (Columnable a, Read a) => T.Text -> Expr (Either T.Text a)-castEither colName = CastWith colName "castEither" (either (Left . T.pack) Right)--{- | Lenient coercion for assertedly non-nullable columns.-Substitutes @error@ for @Nothing@, so it will crash at evaluation time if-any @Nothing@ is actually encountered. For non-nullable and-fully-populated nullable columns no cost is paid.--}-unsafeCast :: forall a. (Columnable a, Read a) => T.Text -> Expr a-unsafeCast colName =- CastWith- colName- "unsafeCast"- (fromRight (error "unsafeCast: unexpected Nothing in column"))--castExpr ::- forall b src.- (Columnable b, Columnable src, Read b) =>- Expr src ->- Expr (Maybe b)-castExpr = CastExprWith @b @(Maybe b) @src "castExpr" (either (const Nothing) Just)--castExprWithDefault ::- forall b src. (Columnable b, Columnable src, Read b) => b -> Expr src -> Expr b-castExprWithDefault def =- CastExprWith @b @b @src- ("castExprWithDefault:" <> T.pack (show def))- (fromRight def)--castExprEither ::- forall b src.- (Columnable b, Columnable src, Read b) =>- Expr src ->- Expr (Either T.Text b)-castExprEither =- CastExprWith @b @(Either T.Text b) @src- "castExprEither"- (either (Left . T.pack) Right)--unsafeCastExpr ::- forall b src. (Columnable b, Columnable src, Read b) => Expr src -> Expr b-unsafeCastExpr =- CastExprWith @b @b @src- "unsafeCastExpr"- (fromRight (error "unsafeCastExpr: unexpected Nothing in column"))--toDouble :: (Columnable a, Real a) => Expr a -> Expr Double-toDouble =- Unary- ( MkUnaryOp- { unaryFn = realToFrac- , unaryName = "toDouble"- , unarySymbol = Nothing- }- )--infix 8 `div`-div :: (Integral a, Columnable a) => Expr a -> Expr a -> Expr a-div = lift2Decorated Prelude.div "div" (Just "//") False 7--mod :: (Integral a, Columnable a) => Expr a -> Expr a -> Expr a-mod = lift2Decorated Prelude.mod "mod" Nothing False 7--eq :: (Columnable a, Eq a, a ~ BaseType a) => Expr a -> Expr a -> Expr Bool-eq = lift2Decorated (==) "eq" (Just "==") True 4--lt :: (Columnable a, Ord a, a ~ BaseType a) => Expr a -> Expr a -> Expr Bool-lt = lift2Decorated (<) "lt" (Just "<") False 4--gt :: (Columnable a, Ord a, a ~ BaseType a) => Expr a -> Expr a -> Expr Bool-gt = lift2Decorated (>) "gt" (Just ">") False 4--leq ::- (Columnable a, Ord a, Eq a, a ~ BaseType a) => Expr a -> Expr a -> Expr Bool-leq = lift2Decorated (<=) "leq" (Just "<=") False 4--geq ::- (Columnable a, Ord a, Eq a, a ~ BaseType a) => Expr a -> Expr a -> Expr Bool-geq = lift2Decorated (>=) "geq" (Just ">=") False 4--and :: Expr Bool -> Expr Bool -> Expr Bool-and = (.&&)--or :: Expr Bool -> Expr Bool -> Expr Bool-or = (.||)--not :: Expr Bool -> Expr Bool-not =- Unary- (MkUnaryOp{unaryFn = Prelude.not, unaryName = "not", unarySymbol = Just "~"})--count :: (Columnable a) => Expr a -> Expr Int-count = Agg (MergeAgg "count" (0 :: Int) (\c _ -> c + 1) (+) id)-{-# SPECIALIZE count :: Expr Double -> Expr Int #-}-{-# SPECIALIZE count :: Expr Float -> Expr Int #-}-{-# SPECIALIZE count :: Expr Int -> Expr Int #-}-{-# SPECIALIZE count :: Expr Int8 -> Expr Int #-}-{-# SPECIALIZE count :: Expr Int16 -> Expr Int #-}-{-# SPECIALIZE count :: Expr Int32 -> Expr Int #-}-{-# SPECIALIZE count :: Expr Int64 -> Expr Int #-}-{-# INLINEABLE count #-}---- | Row count, the equivalent of SQL's @COUNT(*)@.-countAll :: Expr Int-countAll = count (Lit (0 :: Int))-{-# INLINE countAll #-}--collect :: (Columnable a) => Expr a -> Expr [a]-collect = Agg (FoldAgg "collect" (Just []) (flip (:)))-{-# SPECIALIZE collect :: Expr Double -> Expr [Double] #-}-{-# SPECIALIZE collect :: Expr Float -> Expr [Float] #-}-{-# SPECIALIZE collect :: Expr Int -> Expr [Int] #-}-{-# INLINEABLE collect #-}--mode :: (Ord a, Columnable a, Eq a) => Expr a -> Expr a-mode =- Agg- ( CollectAgg- "mode"- ( fst- . L.maximumBy (compare `on` snd)- . M.toList- . V.foldl' (\m e -> M.insertWith (+) e (1 :: Int) m) M.empty- )- )-{-# SPECIALIZE mode :: Expr Double -> Expr Double #-}-{-# SPECIALIZE mode :: Expr Float -> Expr Float #-}-{-# SPECIALIZE mode :: Expr Int -> Expr Int #-}-{-# SPECIALIZE mode :: Expr Int8 -> Expr Int8 #-}-{-# SPECIALIZE mode :: Expr Int16 -> Expr Int16 #-}-{-# SPECIALIZE mode :: Expr Int32 -> Expr Int32 #-}-{-# SPECIALIZE mode :: Expr Int64 -> Expr Int64 #-}-{-# INLINEABLE mode #-}--minimum :: (Columnable a, Ord a) => Expr a -> Expr a-minimum = Agg (FoldAgg "minimum" Nothing Prelude.min)-{-# SPECIALIZE DataFrame.Functions.minimum :: Expr Double -> Expr Double #-}-{-# SPECIALIZE DataFrame.Functions.minimum :: Expr Float -> Expr Float #-}-{-# SPECIALIZE DataFrame.Functions.minimum :: Expr Int -> Expr Int #-}-{-# SPECIALIZE DataFrame.Functions.minimum :: Expr Int8 -> Expr Int8 #-}-{-# SPECIALIZE DataFrame.Functions.minimum :: Expr Int16 -> Expr Int16 #-}-{-# SPECIALIZE DataFrame.Functions.minimum :: Expr Int32 -> Expr Int32 #-}-{-# SPECIALIZE DataFrame.Functions.minimum :: Expr Int64 -> Expr Int64 #-}-{-# INLINEABLE DataFrame.Functions.minimum #-}--maximum :: (Columnable a, Ord a) => Expr a -> Expr a-maximum = Agg (FoldAgg "maximum" Nothing Prelude.max)-{-# SPECIALIZE DataFrame.Functions.maximum :: Expr Double -> Expr Double #-}-{-# SPECIALIZE DataFrame.Functions.maximum :: Expr Float -> Expr Float #-}-{-# SPECIALIZE DataFrame.Functions.maximum :: Expr Int -> Expr Int #-}-{-# SPECIALIZE DataFrame.Functions.maximum :: Expr Int8 -> Expr Int8 #-}-{-# SPECIALIZE DataFrame.Functions.maximum :: Expr Int16 -> Expr Int16 #-}-{-# SPECIALIZE DataFrame.Functions.maximum :: Expr Int32 -> Expr Int32 #-}-{-# SPECIALIZE DataFrame.Functions.maximum :: Expr Int64 -> Expr Int64 #-}-{-# INLINEABLE DataFrame.Functions.maximum #-}--sum :: forall a. (Columnable a, Num a) => Expr a -> Expr a-sum = Agg (FoldAgg "sum" Nothing (+))-{-# SPECIALIZE DataFrame.Functions.sum :: Expr Double -> Expr Double #-}-{-# SPECIALIZE DataFrame.Functions.sum :: Expr Float -> Expr Float #-}-{-# SPECIALIZE DataFrame.Functions.sum :: Expr Int -> Expr Int #-}-{-# SPECIALIZE DataFrame.Functions.sum :: Expr Int8 -> Expr Int8 #-}-{-# SPECIALIZE DataFrame.Functions.sum :: Expr Int16 -> Expr Int16 #-}-{-# SPECIALIZE DataFrame.Functions.sum :: Expr Int32 -> Expr Int32 #-}-{-# SPECIALIZE DataFrame.Functions.sum :: Expr Int64 -> Expr Int64 #-}-{-# INLINEABLE DataFrame.Functions.sum #-}--sumMaybe :: forall a. (Columnable a, Num a) => Expr (Maybe a) -> Expr a-sumMaybe = Agg (CollectAgg "sumMaybe" (P.sum . Maybe.catMaybes . V.toList))-{-# SPECIALIZE sumMaybe :: Expr (Maybe Double) -> Expr Double #-}-{-# SPECIALIZE sumMaybe :: Expr (Maybe Float) -> Expr Float #-}-{-# SPECIALIZE sumMaybe :: Expr (Maybe Int) -> Expr Int #-}-{-# SPECIALIZE sumMaybe :: Expr (Maybe Int8) -> Expr Int8 #-}-{-# SPECIALIZE sumMaybe :: Expr (Maybe Int16) -> Expr Int16 #-}-{-# SPECIALIZE sumMaybe :: Expr (Maybe Int32) -> Expr Int32 #-}-{-# SPECIALIZE sumMaybe :: Expr (Maybe Int64) -> Expr Int64 #-}-{-# INLINEABLE sumMaybe #-}--mean :: (Columnable a, Real a) => Expr a -> Expr Double-mean =- Agg- ( MergeAgg- "mean"- (MeanAcc 0.0 0)- (\(MeanAcc s c) x -> MeanAcc (s + realToFrac x) (c + 1))- (\(MeanAcc s1 c1) (MeanAcc s2 c2) -> MeanAcc (s1 + s2) (c1 + c2))- (\(MeanAcc s c) -> if c == 0 then 0 / 0 else s / fromIntegral c)- )-{-# SPECIALIZE mean :: Expr Double -> Expr Double #-}-{-# SPECIALIZE mean :: Expr Float -> Expr Double #-}-{-# SPECIALIZE mean :: Expr Int -> Expr Double #-}-{-# SPECIALIZE mean :: Expr Int8 -> Expr Double #-}-{-# SPECIALIZE mean :: Expr Int16 -> Expr Double #-}-{-# SPECIALIZE mean :: Expr Int32 -> Expr Double #-}-{-# SPECIALIZE mean :: Expr Int64 -> Expr Double #-}-{-# INLINEABLE mean #-}--meanMaybe :: forall a. (Columnable a, Real a) => Expr (Maybe a) -> Expr Double-meanMaybe = Agg (CollectAgg "meanMaybe" (mean' . optionalToDoubleVector))-{-# SPECIALIZE meanMaybe :: Expr (Maybe Double) -> Expr Double #-}-{-# SPECIALIZE meanMaybe :: Expr (Maybe Float) -> Expr Double #-}-{-# SPECIALIZE meanMaybe :: Expr (Maybe Int) -> Expr Double #-}-{-# SPECIALIZE meanMaybe :: Expr (Maybe Int8) -> Expr Double #-}-{-# SPECIALIZE meanMaybe :: Expr (Maybe Int16) -> Expr Double #-}-{-# SPECIALIZE meanMaybe :: Expr (Maybe Int32) -> Expr Double #-}-{-# SPECIALIZE meanMaybe :: Expr (Maybe Int64) -> Expr Double #-}-{-# INLINEABLE meanMaybe #-}--variance :: (Columnable a, Real a, VU.Unbox a) => Expr a -> Expr Double-variance = Agg (CollectAgg "variance" variance')-{-# SPECIALIZE variance :: Expr Double -> Expr Double #-}-{-# SPECIALIZE variance :: Expr Float -> Expr Double #-}-{-# SPECIALIZE variance :: Expr Int -> Expr Double #-}-{-# SPECIALIZE variance :: Expr Int8 -> Expr Double #-}-{-# SPECIALIZE variance :: Expr Int16 -> Expr Double #-}-{-# SPECIALIZE variance :: Expr Int32 -> Expr Double #-}-{-# SPECIALIZE variance :: Expr Int64 -> Expr Double #-}-{-# INLINEABLE variance #-}--median :: (Columnable a, Real a, VU.Unbox a) => Expr a -> Expr Double-median = Agg (CollectAgg "median" median')-{-# SPECIALIZE median :: Expr Double -> Expr Double #-}-{-# SPECIALIZE median :: Expr Float -> Expr Double #-}-{-# SPECIALIZE median :: Expr Int -> Expr Double #-}-{-# SPECIALIZE median :: Expr Int8 -> Expr Double #-}-{-# SPECIALIZE median :: Expr Int16 -> Expr Double #-}-{-# SPECIALIZE median :: Expr Int32 -> Expr Double #-}-{-# SPECIALIZE median :: Expr Int64 -> Expr Double #-}-{-# INLINEABLE median #-}--medianMaybe :: (Columnable a, Real a) => Expr (Maybe a) -> Expr Double-medianMaybe = Agg (CollectAgg "meanMaybe" (median' . optionalToDoubleVector))-{-# SPECIALIZE medianMaybe :: Expr (Maybe Double) -> Expr Double #-}-{-# SPECIALIZE medianMaybe :: Expr (Maybe Float) -> Expr Double #-}-{-# SPECIALIZE medianMaybe :: Expr (Maybe Int) -> Expr Double #-}-{-# SPECIALIZE medianMaybe :: Expr (Maybe Int8) -> Expr Double #-}-{-# SPECIALIZE medianMaybe :: Expr (Maybe Int16) -> Expr Double #-}-{-# SPECIALIZE medianMaybe :: Expr (Maybe Int32) -> Expr Double #-}-{-# SPECIALIZE medianMaybe :: Expr (Maybe Int64) -> Expr Double #-}-{-# INLINEABLE medianMaybe #-}--optionalToDoubleVector :: (Real a) => V.Vector (Maybe a) -> VU.Vector Double-optionalToDoubleVector =- VU.fromList- . V.foldl'- (\acc e -> if Maybe.isJust e then realToFrac (Maybe.fromMaybe 0 e) : acc else acc)- []--percentile :: Int -> Expr Double -> Expr Double-percentile n =- Agg- ( CollectAgg- (T.pack $ "percentile " ++ show n)- (percentile' n)- )--stddev :: (Columnable a, Real a, VU.Unbox a) => Expr a -> Expr Double-stddev = Agg (CollectAgg "stddev" (sqrt . variance'))-{-# SPECIALIZE stddev :: Expr Double -> Expr Double #-}-{-# SPECIALIZE stddev :: Expr Float -> Expr Double #-}-{-# SPECIALIZE stddev :: Expr Int -> Expr Double #-}-{-# SPECIALIZE stddev :: Expr Int8 -> Expr Double #-}-{-# SPECIALIZE stddev :: Expr Int16 -> Expr Double #-}-{-# SPECIALIZE stddev :: Expr Int32 -> Expr Double #-}-{-# SPECIALIZE stddev :: Expr Int64 -> Expr Double #-}-{-# INLINEABLE stddev #-}--stddevMaybe :: forall a. (Columnable a, Real a) => Expr (Maybe a) -> Expr Double-stddevMaybe = Agg (CollectAgg "stddevMaybe" (sqrt . variance' . optionalToDoubleVector))-{-# SPECIALIZE stddevMaybe :: Expr (Maybe Double) -> Expr Double #-}-{-# SPECIALIZE stddevMaybe :: Expr (Maybe Float) -> Expr Double #-}-{-# SPECIALIZE stddevMaybe :: Expr (Maybe Int) -> Expr Double #-}-{-# SPECIALIZE stddevMaybe :: Expr (Maybe Int8) -> Expr Double #-}-{-# SPECIALIZE stddevMaybe :: Expr (Maybe Int16) -> Expr Double #-}-{-# SPECIALIZE stddevMaybe :: Expr (Maybe Int32) -> Expr Double #-}-{-# SPECIALIZE stddevMaybe :: Expr (Maybe Int64) -> Expr Double #-}-{-# INLINEABLE stddevMaybe #-}--zScore :: Expr Double -> Expr Double-zScore c = (c - mean c) / stddev c--pow :: (Columnable a, Num a) => Expr a -> Int -> Expr a-pow expr i = lift2Decorated (^) "pow" (Just "^") True 8 expr (Lit i)-{-# SPECIALIZE pow :: Expr Double -> Int -> Expr Double #-}-{-# SPECIALIZE pow :: Expr Float -> Int -> Expr Float #-}-{-# SPECIALIZE pow :: Expr Int -> Int -> Expr Int #-}-{-# INLINEABLE pow #-}--relu :: (Columnable a, Num a, Ord a) => Expr a -> Expr a-relu = liftDecorated (Prelude.max 0) "relu" Nothing-{-# SPECIALIZE relu :: Expr Double -> Expr Double #-}-{-# SPECIALIZE relu :: Expr Float -> Expr Float #-}-{-# SPECIALIZE relu :: Expr Int -> Expr Int #-}-{-# INLINEABLE relu #-}--min :: (Columnable a, Ord a) => Expr a -> Expr a -> Expr a-min = lift2Decorated Prelude.min "min" Nothing True 1-{-# SPECIALIZE DataFrame.Functions.min ::- Expr Double -> Expr Double -> Expr Double- #-}-{-# SPECIALIZE DataFrame.Functions.min ::- Expr Float -> Expr Float -> Expr Float- #-}-{-# SPECIALIZE DataFrame.Functions.min :: Expr Int -> Expr Int -> Expr Int #-}-{-# INLINEABLE DataFrame.Functions.min #-}--max :: (Columnable a, Ord a) => Expr a -> Expr a -> Expr a-max = lift2Decorated Prelude.max "max" Nothing True 1-{-# SPECIALIZE DataFrame.Functions.max ::- Expr Double -> Expr Double -> Expr Double- #-}-{-# SPECIALIZE DataFrame.Functions.max ::- Expr Float -> Expr Float -> Expr Float- #-}-{-# SPECIALIZE DataFrame.Functions.max :: Expr Int -> Expr Int -> Expr Int #-}-{-# INLINEABLE DataFrame.Functions.max #-}--reduce ::- forall a b.- (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 #-}--toMaybe :: (Columnable a) => Expr a -> Expr (Maybe a)-toMaybe = liftDecorated Just "toMaybe" Nothing-{-# SPECIALIZE toMaybe :: Expr Double -> Expr (Maybe Double) #-}-{-# SPECIALIZE toMaybe :: Expr Float -> Expr (Maybe Float) #-}-{-# SPECIALIZE toMaybe :: Expr Int -> Expr (Maybe Int) #-}-{-# INLINEABLE toMaybe #-}--fromMaybe :: (Columnable a) => a -> Expr (Maybe a) -> Expr a-fromMaybe d = liftDecorated (Maybe.fromMaybe d) "fromMaybe" Nothing-{-# SPECIALIZE fromMaybe :: Double -> Expr (Maybe Double) -> Expr Double #-}-{-# SPECIALIZE fromMaybe :: Float -> Expr (Maybe Float) -> Expr Float #-}-{-# SPECIALIZE fromMaybe :: Int -> Expr (Maybe Int) -> Expr Int #-}-{-# INLINEABLE fromMaybe #-}--isJust :: (Columnable a) => Expr (Maybe a) -> Expr Bool-isJust = liftDecorated Maybe.isJust "isJust" Nothing-{-# SPECIALIZE isJust :: Expr (Maybe Double) -> Expr Bool #-}-{-# SPECIALIZE isJust :: Expr (Maybe Int) -> Expr Bool #-}-{-# INLINEABLE isJust #-}--isNothing :: (Columnable a) => Expr (Maybe a) -> Expr Bool-isNothing = liftDecorated Maybe.isNothing "isNothing" Nothing-{-# SPECIALIZE isNothing :: Expr (Maybe Double) -> Expr Bool #-}-{-# SPECIALIZE isNothing :: Expr (Maybe Int) -> Expr Bool #-}-{-# INLINEABLE isNothing #-}--fromJust :: (Columnable a) => Expr (Maybe a) -> Expr a-fromJust = liftDecorated Maybe.fromJust "fromJust" Nothing-{-# SPECIALIZE fromJust :: Expr (Maybe Double) -> Expr Double #-}-{-# SPECIALIZE fromJust :: Expr (Maybe Int) -> Expr Int #-}-{-# INLINEABLE fromJust #-}--whenPresent ::- forall a 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)-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)-recode mapping =- Unary- ( MkUnaryOp- { unaryFn = (`lookup` mapping)- , unaryName = "recode " <> T.pack (show mapping)- , unarySymbol = Nothing- }- )--recodeWithCondition ::- forall a b.- (Columnable a, Columnable 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-recodeWithDefault d mapping =- Unary- ( MkUnaryOp- { unaryFn = Maybe.fromMaybe d . (`lookup` mapping)- , unaryName =- "recodeWithDefault " <> T.pack (show d) <> " " <> T.pack (show mapping)- , unarySymbol = Nothing- }- )--firstOrNothing :: (Columnable a) => Expr [a] -> Expr (Maybe a)-firstOrNothing = liftDecorated Maybe.listToMaybe "firstOrNothing" Nothing--lastOrNothing :: (Columnable a) => Expr [a] -> Expr (Maybe a)-lastOrNothing = liftDecorated (Maybe.listToMaybe . reverse) "lastOrNothing" Nothing--splitOn :: T.Text -> Expr T.Text -> Expr [T.Text]-splitOn delim = liftDecorated (T.splitOn delim) "splitOn" Nothing--match :: T.Text -> Expr T.Text -> Expr (Maybe T.Text)-match regex =- liftDecorated- ((\r -> if T.null r then Nothing else Just r) . (=~ regex))- ("match " <> T.pack (show regex))- Nothing--matchAll :: T.Text -> Expr T.Text -> Expr [T.Text]-matchAll regex =- liftDecorated- (getAllTextMatches . (=~ regex))- ("matchAll " <> T.pack (show regex))- Nothing--parseDate ::- (ParseTime t, Columnable t) => T.Text -> Expr T.Text -> Expr (Maybe t)-parseDate format =- liftDecorated- (parseTimeM True defaultTimeLocale (T.unpack format) . T.unpack)- ("parseDate " <> format)- Nothing--daysBetween :: Expr Day -> Expr Day -> Expr Int-daysBetween =- lift2Decorated- (\d1 d2 -> fromIntegral (diffDays d1 d2))- "daysBetween"- Nothing- True- 2--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)-bind f = liftDecorated (>>= f) "bind" Nothing--{- | Window function: evaluate an expression partitioned by the given columns.--Each partition computes the inner expression independently, and the result-is broadcast back to every row in that partition. This is analogous to-Polars' @.over()@ or SQL @OVER (PARTITION BY ...)@.--@--- Per-country median, broadcast to every row:-F.over [\"country\"] (F.median (F.col \@Double \"amount\"))---- Deviation from group mean:-F.col \@Double \"amount\" - F.over [\"group\"] (F.mean (F.col \@Double \"amount\"))-@--}-over :: (Columnable a) => [T.Text] -> Expr a -> Expr a-over = Over---- See Section 2.4 of the Haskell Report https://www.haskell.org/definition/haskell2010.pdf-isReservedId :: T.Text -> Bool-isReservedId t = case t of- "case" -> True- "class" -> True- "data" -> True- "default" -> True- "deriving" -> True- "do" -> True- "else" -> True- "foreign" -> True- "if" -> True- "import" -> True- "in" -> True- "infix" -> True- "infixl" -> True- "infixr" -> True- "instance" -> True- "let" -> True- "module" -> True- "newtype" -> True- "of" -> True- "then" -> True- "type" -> True- "where" -> True- _ -> False--isVarId :: T.Text -> Bool-isVarId t = case T.uncons t of- -- We might want to check c == '_' || Char.isLower c- -- since the haskell report considers '_' a lowercase character- -- However, to prevent an edge case where a user may have a- -- "Name" and an "_Name_" in the same scope, wherein we'd end up- -- with duplicate "_Name_"s, we eschew the check for '_' here.- Just (c, _) -> Char.isLower c && Char.isAlpha c- Nothing -> False--isHaskellIdentifier :: T.Text -> Bool-isHaskellIdentifier t = Prelude.not (isVarId t) || isReservedId t--sanitize :: T.Text -> T.Text-sanitize t- | isValid = t- | isHaskellIdentifier t' = "_" <> t' <> "_"- | otherwise = t'- where- isValid =- Prelude.not (isHaskellIdentifier t)- && isVarId t- && T.all Char.isAlphaNum t- t' = T.map replaceInvalidCharacters . T.filter (Prelude.not . parentheses) $ t- replaceInvalidCharacters c- | Char.isUpper c = Char.toLower c- | Char.isSpace c = '_'- | Char.isPunctuation c = '_' -- '-' will also become a '_'- | Char.isSymbol c = '_'- | Char.isAlphaNum c = c -- Blanket condition- | otherwise = '_' -- If we're unsure we'll default to an underscore- parentheses c = case c of- '(' -> True- ')' -> True- '{' -> True- '}' -> True- '[' -> True- ']' -> True- _ -> False
− src/DataFrame/IO/CSV.hs
@@ -1,794 +0,0 @@-{-# LANGUAGE BangPatterns #-}-{-# LANGUAGE ExplicitNamespaces #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE NumericUnderscores #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE RankNTypes #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}--module DataFrame.IO.CSV where--import qualified Data.ByteString as BS-import qualified Data.ByteString.Char8 as C-import qualified Data.ByteString.Lazy as BL-import qualified Data.Map.Strict as M-import qualified Data.Proxy as P-import qualified Data.Text as T-import qualified Data.Text.Encoding as TE-import qualified Data.Text.IO as TIO-import qualified Data.Vector as V-import qualified Data.Vector.Mutable as VM-import qualified Data.Vector.Unboxed as VU-import qualified Data.Vector.Unboxed.Mutable as VUM--import Data.Csv.Streaming (Records (..))-import qualified Data.Csv.Streaming as CsvStream--import Control.DeepSeq-import Control.Exception (SomeException, catch)-import Control.Monad-import Control.Monad.ST (runST)-import Data.Char-import qualified Data.Csv as Csv-import Data.Either-import Data.Functor-import Data.IORef-import Data.Maybe-import Data.Type.Equality (TestEquality (testEquality))-import Data.Word (Word8)-import DataFrame.Internal.Column-import DataFrame.Internal.DataFrame (DataFrame (..), toSeparated)-import DataFrame.Internal.Parsing-import DataFrame.Internal.Schema-import DataFrame.Operations.Typing-import System.IO-import Type.Reflection-import Prelude hiding (concat, takeWhile)--chunkSize :: Int-chunkSize = 16_384--data PagedVector a = PagedVector- { pvChunks :: !(IORef [V.Vector a])- -- ^ Finished chunks (reverse order)- , pvActive :: !(IORef (VM.IOVector a))- -- ^ Current mutable chunk- , pvCount :: !(IORef Int)- -- ^ Items written in current chunk- }--data PagedUnboxedVector a = PagedUnboxedVector- { puvChunks :: !(IORef [VU.Vector a])- , puvActive :: !(IORef (VUM.IOVector a))- , puvCount :: !(IORef Int)- }--data BuilderColumn- = BuilderInt !(PagedUnboxedVector Int) !(PagedUnboxedVector Word8)- | BuilderDouble !(PagedUnboxedVector Double) !(PagedUnboxedVector Word8)- | BuilderText !(PagedVector T.Text) !(PagedUnboxedVector Word8)- | BuilderBS !(PagedVector BS.ByteString) !(PagedUnboxedVector Word8)--newPagedVector :: IO (PagedVector a)-newPagedVector = do- active <- VM.unsafeNew chunkSize- PagedVector <$> newIORef [] <*> newIORef active <*> newIORef 0--newPagedUnboxedVector :: (VUM.Unbox a) => IO (PagedUnboxedVector a)-newPagedUnboxedVector = do- active <- VUM.unsafeNew chunkSize- PagedUnboxedVector <$> newIORef [] <*> newIORef active <*> newIORef 0--appendPagedVector :: PagedVector a -> a -> IO ()-appendPagedVector (PagedVector chunksRef activeRef countRef) !val = do- count <- readIORef countRef- active <- readIORef activeRef-- if count < chunkSize- then do- VM.unsafeWrite active count val- writeIORef countRef $! count + 1- else do- frozen <- V.unsafeFreeze active- modifyIORef' chunksRef (frozen :)-- newActive <- VM.unsafeNew chunkSize- VM.unsafeWrite newActive 0 val-- writeIORef activeRef newActive- writeIORef countRef 1-{-# INLINE appendPagedVector #-}--appendPagedUnboxedVector :: (VUM.Unbox a) => PagedUnboxedVector a -> a -> IO ()-appendPagedUnboxedVector (PagedUnboxedVector chunksRef activeRef countRef) !val = do- count <- readIORef countRef- active <- readIORef activeRef-- if count < chunkSize- then do- VUM.unsafeWrite active count val- writeIORef countRef $! count + 1- else do- frozen <- VU.unsafeFreeze active- modifyIORef' chunksRef (frozen :)-- newActive <- VUM.unsafeNew chunkSize- VUM.unsafeWrite newActive 0 val-- writeIORef activeRef newActive- writeIORef countRef 1-{-# INLINE appendPagedUnboxedVector #-}--freezePagedVector :: PagedVector a -> IO (V.Vector a)-freezePagedVector (PagedVector chunksRef activeRef countRef) = do- count <- readIORef countRef- active <- readIORef activeRef- chunks <- readIORef chunksRef-- writeIORef chunksRef [] -- release chunk references- let frozenChunks = reverse chunks- totalLen = count + sum (map V.length frozenChunks)-- mv <- VM.unsafeNew totalLen-- let copyChunk !offset chunk = do- V.copy (VM.slice offset (V.length chunk) mv) chunk- pure (offset + V.length chunk)-- offset <- foldM copyChunk 0 frozenChunks- VM.copy (VM.slice offset count mv) (VM.slice 0 count active)-- V.unsafeFreeze mv--freezePagedUnboxedVector ::- (VUM.Unbox a) => PagedUnboxedVector a -> IO (VU.Vector a)-freezePagedUnboxedVector (PagedUnboxedVector chunksRef activeRef countRef) = do- count <- readIORef countRef- active <- readIORef activeRef- chunks <- readIORef chunksRef-- writeIORef chunksRef [] -- release chunk references- let frozenChunks = reverse chunks- totalLen = count + sum (map VU.length frozenChunks)-- mv <- VUM.unsafeNew totalLen-- let copyChunk !offset chunk = do- VU.copy (VUM.slice offset (VU.length chunk) mv) chunk- pure (offset + VU.length chunk)-- offset <- foldM copyChunk 0 frozenChunks- VUM.copy (VUM.slice offset count mv) (VUM.slice 0 count active)-- VU.unsafeFreeze mv---- | STANDARD CONFIG TYPES-data HeaderSpec = NoHeader | UseFirstRow | ProvideNames [T.Text]- deriving (Eq, Show)--data TypeSpec- = InferFromSample Int- | SpecifyTypes [(T.Text, SchemaType)] TypeSpec- | NoInference--{- | How the fast reader should treat a row whose field count does not-match the header row. Only consulted by @dataframe-fastcsv@; the pure-Haskell reader has its own semantics.--}-data RaggedRowPolicy- = {- | Fill missing cells with nulls; silently drop extras.- Matches pandas / polars lenient defaults.- -}- PadWithNull- | {- | Fill missing cells with nulls; silently drop extras. Alias- kept for ergonomic naming when the caller only cares about the- "don't raise, just forget the extras" half of 'PadWithNull'.- -}- Truncate- | {- | Raise a 'DataFrame.IO.CSV.Fast.CsvParseError' on any row whose- field count differs from the header. Matches polars' strict- schema-bound mode.- -}- RaiseOnRagged- deriving (Eq, Show)---- | How the fast reader should treat an unclosed quoted field at EOF.-data UnclosedQuotePolicy- = -- | Raise a 'DataFrame.IO.CSV.Fast.CsvParseError'. Default.- RaiseOnUnclosedQuote- | {- | Return whatever rows were parsed before the stray quote; the- remainder is silently dropped.- -}- BestEffort- deriving (Eq, Show)---- | CSV read parameters.-data ReadOptions = ReadOptions- { headerSpec :: HeaderSpec- -- ^ Where to get the headers from. (default: UseFirstRow)- , typeSpec :: TypeSpec- -- ^ Whether/how to infer types. (default: InferFromSample 100)- , safeRead :: SafeReadMode- {- ^ Default 'SafeReadMode' for columns without an entry in- 'safeReadOverrides'. (default: 'NoSafeRead')- -}- , safeReadOverrides :: [(T.Text, SafeReadMode)]- -- ^ Per-column 'SafeReadMode' overrides; takes precedence over 'safeRead'.- , dateFormat :: String- {- ^ Format of date fields as recognized by the Data.Time.Format module.-- __Examples:__-- @- > parseTimeM True defaultTimeLocale "%Y/%-m/%-d" "2010/3/04" :: Maybe Day- Just 2010-03-04- > parseTimeM True defaultTimeLocale "%d/%-m/%-Y" "04/3/2010" :: Maybe Day- Just 2010-03-04- @- -}- , columnSeparator :: Char- -- ^ Character that separates column values.- , numColumns :: Maybe Int- -- ^ Number of columns to read.- , missingIndicators :: [T.Text]- -- ^ Values that should be read as `Nothing`.- , fastCsvOnRaggedRow :: RaggedRowPolicy- {- ^ @dataframe-fastcsv@: how to treat rows with a non-header field count.- (default: 'PadWithNull')- -}- , fastCsvOnUnclosedQuote :: UnclosedQuotePolicy- {- ^ @dataframe-fastcsv@: how to treat an unclosed quoted field at EOF.- (default: 'RaiseOnUnclosedQuote')- -}- , fastCsvTrimUnquoted :: Bool- {- ^ @dataframe-fastcsv@: if 'True', leading/trailing whitespace is- stripped from unquoted fields after decoding. RFC 4180 preserves- this whitespace, and that is the default ('False').- -}- }--shouldInferFromSample :: TypeSpec -> Bool-shouldInferFromSample (InferFromSample _) = True-shouldInferFromSample (SpecifyTypes _ fallback) = shouldInferFromSample fallback-shouldInferFromSample _ = False--schemaTypeMap :: TypeSpec -> M.Map T.Text SchemaType-schemaTypeMap (SpecifyTypes xs _) = M.fromList xs-schemaTypeMap _ = M.empty--typeInferenceSampleSize :: TypeSpec -> Int-typeInferenceSampleSize (InferFromSample n) = n-typeInferenceSampleSize (SpecifyTypes _ fallback) = typeInferenceSampleSize fallback-typeInferenceSampleSize _ = 0--defaultReadOptions :: ReadOptions-defaultReadOptions =- ReadOptions- { headerSpec = UseFirstRow- , typeSpec = InferFromSample 100- , safeRead = NoSafeRead- , safeReadOverrides = []- , dateFormat = "%Y-%m-%d"- , columnSeparator = ','- , numColumns = Nothing- , missingIndicators =- ["Nothing", "NULL", "", " ", "nan", "null", "N/A", "NaN", "NAN", "NA"]- , fastCsvOnRaggedRow = PadWithNull- , fastCsvOnUnclosedQuote = RaiseOnUnclosedQuote- , fastCsvTrimUnquoted = False- }--{- | Read CSV file from path and load it into a dataframe.--==== __Example__-@-ghci> D.readCsv ".\/data\/taxi.csv"--@--}-readCsv :: FilePath -> IO DataFrame-readCsv = readSeparated defaultReadOptions--type CsvReader = Schema -> FilePath -> IO DataFrame--{- | Schema-driven attoparsec CSV reader. Coerces each column to the-type declared in 'Schema'; columns absent from the schema fall back to-the default inference path. Defined in terms of 'readSeparated' with-the 'TypeSpec' filled in.--@-import qualified DataFrame as D-df <- D.readCsvWithSchema schema "input.csv"-@--}-readCsvWithSchema :: CsvReader-readCsvWithSchema schema =- readSeparated- defaultReadOptions- { typeSpec =- SpecifyTypes- (M.toList (elements schema))- (typeSpec defaultReadOptions)- }--{- | Read CSV file from path and load it into a dataframe.--==== __Example__-@-ghci> D.readCsvWithOpts ".\/data\/taxi.csv" (D.defaultReadOptions { dateFormat = "%d/%-m/%-Y" })--@--}-readCsvWithOpts :: ReadOptions -> FilePath -> IO DataFrame-readCsvWithOpts = readSeparated--{- | Read TSV (tab separated) file from path and load it into a dataframe.--==== __Example__-@-ghci> D.readTsv ".\/data\/taxi.tsv"--@--}-readTsv :: FilePath -> IO DataFrame-readTsv = readSeparated (defaultReadOptions{columnSeparator = '\t'})--{- | Read text file with specified delimiter into a dataframe.--==== __Example__-@-ghci> D.readSeparated (D.defaultReadOptions { columnSeparator = ';' }) ".\/data\/taxi.txt"--@--}-readSeparated :: ReadOptions -> FilePath -> IO DataFrame-readSeparated opts !path = do- let stripUtf8Bom bs = fromMaybe bs (BL.stripPrefix "\xEF\xBB\xBF" bs)- csvData <- stripUtf8Bom <$> BL.readFile path- fmap force (decodeSeparated opts csvData)--decodeSeparated :: ReadOptions -> BL.ByteString -> IO DataFrame-decodeSeparated !opts csvData = do- let sep = columnSeparator opts- let decodeOpts = Csv.defaultDecodeOptions{Csv.decDelimiter = fromIntegral (ord sep)}- let stream = CsvStream.decodeWith decodeOpts Csv.NoHeader csvData-- let peekStream (Cons (Right row) rest) = return (row, rest)- peekStream (Cons (Left err) _) = error $ "Error parsing CSV header: " ++ err- peekStream (Nil Nothing _) = error "Empty CSV file"- peekStream (Nil (Just err) _) = error err-- (firstRowRaw, dataStream) <- peekStream stream-- let (columnNames, rowsToProcess) = case headerSpec opts of- NoHeader ->- ( map (T.pack . show) [0 .. V.length firstRowRaw - 1]- , Cons (Right firstRowRaw) dataStream- )- UseFirstRow ->- ( map (T.strip . TE.decodeUtf8Lenient . BL.toStrict) (V.toList firstRowRaw)- , dataStream- )- ProvideNames ns ->- ( ns ++ drop (length ns) (map (T.pack . show) [0 .. V.length firstRowRaw - 1])- , Cons (Right firstRowRaw) dataStream- )-- (sampleRow, _) <- peekStream rowsToProcess- builderCols <- initializeColumns columnNames (V.toList sampleRow) opts- let !builderColsV = V.fromList builderCols- let colNamesV = V.fromList columnNames- resolveMode =- effectiveSafeRead- (safeRead opts)- (safeReadOverrides opts)- -- If ANY column is EitherRead we keep every raw cell (including- -- "N/A" etc.) verbatim; otherwise the missing-indicator list applies.- anyEither =- any (\n -> resolveMode n == EitherRead) columnNames- missing = if anyEither then [] else missingIndicators opts- processStream missing rowsToProcess builderColsV (numColumns opts)-- frozenCols <-- V.zipWithM- (\name bc -> finalizeBuilderColumn (resolveMode name) opts bc)- colNamesV- builderColsV- let numRows = maybe 0 columnLength (frozenCols V.!? 0)-- let df =- DataFrame- frozenCols- (M.fromList (zip columnNames [0 ..]))- (numRows, V.length frozenCols)- M.empty -- TODO give typed column references- pure $ parseWithTypes resolveMode (schemaTypeMap (typeSpec opts)) df--initializeColumns ::- [T.Text] -> [BL.ByteString] -> ReadOptions -> IO [BuilderColumn]-initializeColumns names _row opts = zipWithM initColumn names (map lookupType names)- where- typeMap = schemaTypeMap (typeSpec opts)- -- Return Nothing for columns that should be inferred from BS- shouldInfer = case typeSpec opts of- InferFromSample _ -> True- SpecifyTypes _ fallback -> shouldInferFromSample fallback- NoInference -> False- lookupType name = M.lookup name typeMap- resolveMode =- effectiveSafeRead (safeRead opts) (safeReadOverrides opts)- initColumn :: T.Text -> Maybe SchemaType -> IO BuilderColumn- initColumn name _ | resolveMode name == EitherRead = do- validityRef <- newPagedUnboxedVector- BuilderBS <$> newPagedVector <*> pure validityRef- initColumn _ Nothing | shouldInfer = do- validityRef <- newPagedUnboxedVector- BuilderBS <$> newPagedVector <*> pure validityRef- initColumn _ mtype = do- validityRef <- newPagedUnboxedVector- let t = fromMaybe (schemaType @T.Text) mtype- case t of- SType (_ :: P.Proxy a) -> case testEquality (typeRep @a) (typeRep @Int) of- Just Refl -> BuilderInt <$> newPagedUnboxedVector <*> pure validityRef- Nothing -> case testEquality (typeRep @a) (typeRep @Double) of- Just Refl -> BuilderDouble <$> newPagedUnboxedVector <*> pure validityRef- Nothing -> BuilderText <$> newPagedVector <*> pure validityRef--processStream ::- [T.Text] ->- CsvStream.Records (V.Vector BL.ByteString) ->- V.Vector BuilderColumn ->- Maybe Int ->- IO ()-processStream _ _ _ (Just 0) = return ()-processStream missing (Cons (Right row) rest) cols n =- processRow missing row cols- >> processStream missing rest cols (fmap (flip (-) 1) n)-processStream _missing (Cons (Left err) _) _ _ = error ("CSV Parse Error: " ++ err)-processStream _missing (Nil _ _) _ _ = return ()--processRow ::- [T.Text] -> V.Vector BL.ByteString -> V.Vector BuilderColumn -> IO ()-processRow missing !vals !cols = V.zipWithM_ processValue vals cols- where- processValue !bs !col = do- let !bs' = BL.toStrict bs- case col of- BuilderInt gv valid -> case readByteStringInt bs' of- Just !i -> appendPagedUnboxedVector gv i >> appendPagedUnboxedVector valid 1- Nothing -> appendPagedUnboxedVector gv 0 >> appendPagedUnboxedVector valid 0- BuilderDouble gv valid -> case readByteStringDouble bs' of- Just !d -> appendPagedUnboxedVector gv d >> appendPagedUnboxedVector valid 1- Nothing -> appendPagedUnboxedVector gv 0.0 >> appendPagedUnboxedVector valid 0- BuilderText gv valid -> do- let !val = T.strip (TE.decodeUtf8Lenient bs')- appendPagedVector gv val- appendPagedUnboxedVector valid (if val `elem` missing then 0 else 1)- BuilderBS gv valid -> do- let !bs'' = C.strip bs'- appendPagedVector gv bs''- appendPagedUnboxedVector- valid- (if TE.decodeUtf8Lenient bs'' `elem` missing then 0 else 1)--freezeBuilderColumn :: BuilderColumn -> IO Column-freezeBuilderColumn (BuilderInt gv validRef) = do- vec <- freezePagedUnboxedVector gv- valid <- freezePagedUnboxedVector validRef- if VU.all (== 1) valid- then return $! UnboxedColumn Nothing vec- else constructOptional vec valid-freezeBuilderColumn (BuilderDouble gv validRef) = do- vec <- freezePagedUnboxedVector gv- valid <- freezePagedUnboxedVector validRef- if VU.all (== 1) valid- then return $! UnboxedColumn Nothing vec- else constructOptional vec valid-freezeBuilderColumn (BuilderText gv validRef) = do- vec <- freezePagedVector gv- valid <- freezePagedUnboxedVector validRef- if VU.all (== 1) valid- then return $! BoxedColumn Nothing vec- else constructOptionalBoxed vec valid-freezeBuilderColumn (BuilderBS _ _) =- error- "freezeBuilderColumn: BuilderBS must be finalized via finalizeBuilderColumn"--finalizeBuilderColumn ::- SafeReadMode -> ReadOptions -> BuilderColumn -> IO Column-finalizeBuilderColumn mode opts bc = do- col <- case bc of- BuilderBS gv validRef -> do- vec <- freezePagedVector gv- valid <- freezePagedUnboxedVector validRef- return $! inferColumnFromBS mode opts vec valid- _ -> freezeBuilderColumn bc- return $! case mode of- NoSafeRead -> col- MaybeRead -> ensureOptional col- EitherRead -> col--inferColumnFromBS ::- SafeReadMode ->- ReadOptions ->- V.Vector BS.ByteString ->- VU.Vector Word8 ->- Column-inferColumnFromBS mode opts vec valid =- let sampleN = let n = typeInferenceSampleSize (typeSpec opts) in if n == 0 then 100 else n- dfmt = dateFormat opts- -- The sample IS still Maybe-wrapped; it's bounded at 100 rows- -- by default, so the allocation is ignorable. The previous- -- full-column `asMaybeFull = V.generate ...` allocation is- -- gone — handlers walk (vec, valid) directly.- samples = V.generate (min sampleN (V.length vec)) $ \i ->- if valid VU.! i == 1 then Just (vec V.! i) else Nothing- assumption = makeParsingAssumptionBS dfmt samples- in case mode of- EitherRead -> handleBSEither dfmt assumption vec valid- _ -> case assumption of- IntAssumption -> handleBSInt dfmt vec valid- DoubleAssumption -> handleBSDouble vec valid- BoolAssumption -> handleBSBool vec valid- DateAssumption -> handleBSDate dfmt vec valid- TextAssumption -> handleBSText vec valid- NoAssumption -> handleBSNo dfmt vec valid--{- | 'EitherRead' wrap for the ByteString inference path: produce an-@Either Text a@ column. Raw input bytes are preserved verbatim; rows that were-marked invalid by the builder (e.g. empty cells before EitherRead disabled-missing-indicator detection) become @Left \"\"@.--}-handleBSEither ::- String ->- ParsingAssumption ->- V.Vector BS.ByteString ->- VU.Vector Word8 ->- Column-handleBSEither dfmt assumption vec valid = case assumption of- BoolAssumption -> wrap readByteStringBool- IntAssumption -> wrap readByteStringInt- DoubleAssumption -> wrap readByteStringDouble- DateAssumption -> wrap (readByteStringDate dfmt)- -- Text / No assumption: column is Either Text Text; empty cells become- -- Left "" to keep the "Left means missing/failure" convention.- TextAssumption -> fromVector (V.imap textEither vec)- NoAssumption -> fromVector (V.imap textEither vec)- where- wrap ::- forall a. (Columnable a) => (BS.ByteString -> Maybe a) -> Column- wrap p = fromVector (V.imap (toEither p) vec)-- toEither ::- forall a.- (BS.ByteString -> Maybe a) ->- Int ->- BS.ByteString ->- Either T.Text a- toEither p i bs- | valid VU.! i == 0 = Left (TE.decodeUtf8Lenient bs)- | otherwise = case p bs of- Just v -> Right v- Nothing -> Left (TE.decodeUtf8Lenient bs)-- textEither :: Int -> BS.ByteString -> Either T.Text T.Text- textEither i bs =- let t = TE.decodeUtf8Lenient bs- in if valid VU.! i == 0 || T.null t then Left t else Right t--makeParsingAssumptionBS ::- String -> V.Vector (Maybe BS.ByteString) -> ParsingAssumption-makeParsingAssumptionBS dfmt asMaybe- | V.all (== Nothing) asMaybe = NoAssumption- | vecSameConstructor asMaybe asMaybeBool = BoolAssumption- | vecSameConstructor asMaybe asMaybeInt- && vecSameConstructor asMaybe asMaybeDouble =- IntAssumption- | vecSameConstructor asMaybe asMaybeDouble = DoubleAssumption- | vecSameConstructor asMaybe asMaybeDate = DateAssumption- | otherwise = TextAssumption- where- asMaybeBool = V.map (>>= readByteStringBool) asMaybe- asMaybeInt = V.map (>>= readByteStringInt) asMaybe- asMaybeDouble = V.map (>>= readByteStringDouble) asMaybe- asMaybeDate = V.map (>>= readByteStringDate dfmt) asMaybe---- All @handleBS*@ helpers now take the raw @V.Vector BS.ByteString@--- plus the builder's @VU.Vector Word8@ validity vector, fusing the--- parse + validity check into a single pass via--- 'parseUnboxedColumnWithValid'. The previous @V.Vector (Maybe--- BS.ByteString)@ intermediate — allocated upstream in--- 'inferColumnFromBS' — is gone, along with the paired Int/Double--- parses and the two 'V.zipWith' Bool vectors from--- 'vecSameConstructor'.--handleBSBool ::- V.Vector BS.ByteString -> VU.Vector Word8 -> Column-handleBSBool vec valid =- case parseUnboxedColumnWithValid False readByteStringBool vec valid of- Just (mbm, out) -> UnboxedColumn mbm out- Nothing -> handleBSText vec valid--handleBSInt ::- String -> V.Vector BS.ByteString -> VU.Vector Word8 -> Column-handleBSInt _dfmt vec valid =- case parseUnboxedColumnWithValid 0 readByteStringInt vec valid of- Just (mbm, out) -> UnboxedColumn mbm out- Nothing -> case parseUnboxedColumnWithValid 0 readByteStringDouble vec valid of- Just (mbm, out) -> UnboxedColumn mbm out- Nothing -> handleBSText vec valid--handleBSDouble ::- V.Vector BS.ByteString -> VU.Vector Word8 -> Column-handleBSDouble vec valid =- case parseUnboxedColumnWithValid 0 readByteStringDouble vec valid of- Just (mbm, out) -> UnboxedColumn mbm out- Nothing -> handleBSText vec valid---- Dates are boxed ('Day' isn't 'VU.Unbox'), so fuse into a V.Vector--- (Maybe Day) in one pass. Bails on the first non-null cell that--- fails to parse.-handleBSDate ::- String -> V.Vector BS.ByteString -> VU.Vector Word8 -> Column-handleBSDate dfmt vec valid =- case parseBoxedMaybeBSColumn valid (readByteStringDate dfmt) vec of- Just (anyNull, out)- | anyNull -> fromVector out- | otherwise -> fromVector (V.mapMaybe id out)- Nothing -> handleBSText vec valid---- Fused Text handler: decode each cell's UTF-8 bytes once, mark nulls--- directly from the validity vector. Replaces the two-pass--- `V.map (fmap decodeUtf8Lenient) ... sequenceA ...` pattern.-handleBSText ::- V.Vector BS.ByteString -> VU.Vector Word8 -> Column-handleBSText vec valid- | VU.any (== 0) valid =- fromVector- ( V.imap- ( \i bs ->- if valid VU.! i == 0- then Nothing- else Just (TE.decodeUtf8Lenient bs)- )- vec- )- | otherwise = fromVector (V.map TE.decodeUtf8Lenient vec)--handleBSNo ::- String -> V.Vector BS.ByteString -> VU.Vector Word8 -> Column-handleBSNo dfmt vec valid- | VU.all (== 0) valid =- fromVector (V.map (const (Nothing :: Maybe T.Text)) vec)- | Just (mbm, out) <-- parseUnboxedColumnWithValid False readByteStringBool vec valid =- UnboxedColumn mbm out- | Just (mbm, out) <- parseUnboxedColumnWithValid 0 readByteStringInt vec valid =- UnboxedColumn mbm out- | Just (mbm, out) <- parseUnboxedColumnWithValid 0 readByteStringDouble vec valid =- UnboxedColumn mbm out- | otherwise = case parseBoxedMaybeBSColumn valid (readByteStringDate dfmt) vec of- Just (anyNull, out)- | anyNull -> fromVector out- | otherwise -> fromVector (V.mapMaybe id out)- Nothing -> handleBSText vec valid---- Boxed counterpart to 'parseUnboxedColumnWithValid' for types that--- aren't 'VU.Unbox' (e.g. 'Day'). Same one-pass + early-bail shape.-parseBoxedMaybeBSColumn ::- VU.Vector Word8 ->- (BS.ByteString -> Maybe a) ->- V.Vector BS.ByteString ->- Maybe (Bool, V.Vector (Maybe a))-parseBoxedMaybeBSColumn valid parser vec = runST $ do- let n = V.length vec- out <- VM.new n- let loop !i !anyNull- | i >= n = do- frozen <- V.unsafeFreeze out- return (Just (anyNull, frozen))- | VU.unsafeIndex valid i == 0 = do- VM.unsafeWrite out i Nothing- loop (i + 1) True- | otherwise = case parser (V.unsafeIndex vec i) of- Just v -> do- VM.unsafeWrite out i (Just v)- loop (i + 1) anyNull- Nothing -> return Nothing- loop 0 False--{- | One-pass fused parse into a typed unboxed column. Avoids the-@V.Vector (Maybe a)@ intermediate that the "parse then 'sequenceA'"-idiom requires, and avoids the upstream @V.Vector (Maybe src)@-classification by reading nullability from a precomputed validity-vector (produced by the CSV builder alongside the raw cells).--Returns @Just (mbm, vec)@ only when every non-null cell parses. The-first unparseable cell short-circuits to @Nothing@ so the caller can-fall back to the next assumption (Int → Double → Text). @mbm@ is-@Nothing@ when no nulls exist (the column is non-nullable) and-@Just bm@ otherwise.--Null slots are filled with @nullValue@; downstream consumers only see-them through the bitmap, so the sentinel never escapes.--Memory shape for a length-@n@ input:-- * 1 × 'VUM.STVector' of @a@ (final data, @sizeOf a × n@ bytes)- * 1 × 'VUM.STVector' of 'Word8' (per-element validity, @n@ bytes)- * 1 × 'Bitmap' (bit-packed, @⌈n\/8⌉@ bytes) — only when nulls exist.--}-parseUnboxedColumnWithValid ::- forall src a.- (VU.Unbox a) =>- a ->- (src -> Maybe a) ->- V.Vector src ->- VU.Vector Word8 ->- Maybe (Maybe Bitmap, VU.Vector a)-parseUnboxedColumnWithValid nullValue parser vec valid = runST $ do- let n = V.length vec- values <- VUM.unsafeNew n- vmask <- VUM.unsafeNew n- let go !i !anyNull- | i >= n = finalizeParseResult values vmask anyNull- | VU.unsafeIndex valid i == 0 = do- VUM.unsafeWrite vmask i 0- VUM.unsafeWrite values i nullValue- go (i + 1) True- | otherwise = case parser (V.unsafeIndex vec i) of- Just v -> do- VUM.unsafeWrite vmask i 1- VUM.unsafeWrite values i v- go (i + 1) anyNull- Nothing -> return Nothing- go 0 False-{-# INLINE parseUnboxedColumnWithValid #-}--constructOptional ::- (VU.Unbox a, Columnable a) => VU.Vector a -> VU.Vector Word8 -> IO Column-constructOptional vec valid = do- let bm = buildBitmapFromValid valid- pure $ UnboxedColumn (Just bm) vec--constructOptionalBoxed :: V.Vector T.Text -> VU.Vector Word8 -> IO Column-constructOptionalBoxed vec valid = do- let bm = buildBitmapFromValid valid- pure $ BoxedColumn (Just bm) vec--writeCsv :: FilePath -> DataFrame -> IO ()-writeCsv = writeSeparated ','--writeTsv :: FilePath -> DataFrame -> IO ()-writeTsv = writeSeparated '\t'--writeSeparated ::- -- | Separator- Char ->- -- | Path to write to- FilePath ->- DataFrame ->- IO ()-writeSeparated c filepath df = TIO.writeFile filepath (toSeparated c df)---- | Parse a CSV string into a DataFrame using default options.-fromCsv :: String -> IO (Either String DataFrame)-fromCsv s = do- let bs = BL.fromStrict (TE.encodeUtf8 (T.pack s))- (Right <$> decodeSeparated defaultReadOptions bs)- `catch` (\(e :: SomeException) -> pure (Left (show e)))---- | Parse a lazy 'ByteString' containing CSV data into a DataFrame using default options.-fromCsvBytes :: BL.ByteString -> IO DataFrame-fromCsvBytes = decodeSeparated defaultReadOptions--stripQuotes :: T.Text -> T.Text-stripQuotes txt =- case T.uncons txt of- Just ('"', rest) ->- case T.unsnoc rest of- Just (middle, '"') -> middle- _ -> txt- _ -> txt
− src/DataFrame/IO/JSON.hs
@@ -1,133 +0,0 @@-{-# LANGUAGE TypeApplications #-}--module DataFrame.IO.JSON (- readJSON,- readJSONEither,-) where--import Control.Monad (forM)-import Data.Aeson-import qualified Data.Aeson.Key as K-import qualified Data.Aeson.KeyMap as KM-import qualified Data.ByteString.Lazy as LBS-import Data.Maybe (catMaybes)-import Data.Scientific (toRealFloat)-import Data.Text (Text)-import qualified Data.Text as T-import qualified Data.Vector as V--import qualified DataFrame.Internal.Column as D-import qualified DataFrame.Internal.DataFrame as D-import qualified DataFrame.Operations.Core as D--readJSONEither :: LBS.ByteString -> Either String D.DataFrame-readJSONEither bs = do- v <- note "Could not decode JSON" (decode @Value bs)- rows <- toArrayOfObjects v- let cols :: [Text]- cols =- uniq- . concatMap (map K.toText . KM.keys)- . V.toList- $ rows-- columns <- forM cols $ \c -> do- let col = buildColumn rows c- pure (c, col)-- pure $ D.fromNamedColumns columns--readJSON :: FilePath -> IO D.DataFrame-readJSON path = do- contents <- LBS.readFile path- case readJSONEither contents of- Left err -> fail $ "readJSON: " <> err- Right df -> pure df--toArrayOfObjects :: Value -> Either String (V.Vector Object)-toArrayOfObjects (Array xs)- | V.null xs = Left "Top-level JSON array is empty"- | otherwise = traverse asObject xs-toArrayOfObjects _ =- Left "Top-level JSON value must be a JSON array of objects"--asObject :: Value -> Either String Object-asObject (Object o) = Right o-asObject _ = Left "Expected each element of the array to be an object"--uniq :: (Ord a) => [a] -> [a]-uniq = go mempty- where- go _ [] = []- go seen (x : xs)- | x `elem` seen = go seen xs- | otherwise = x : go (x : seen) xs--note :: e -> Maybe a -> Either e a-note e = maybe (Left e) Right--data ColType- = CTString- | CTNumber- | CTBool- | CTArray- | CTMixed--buildColumn :: V.Vector Object -> Text -> D.Column-buildColumn rows colName =- let key = K.fromText colName- values :: V.Vector (Maybe Value)- values = V.map (KM.lookup key) rows- colType = detectColType values- in case colType of- CTString ->- D.fromVector (fmap (fmap asText) values)- CTNumber ->- D.fromVector (fmap (fmap asDouble) values)- CTBool ->- D.fromVector (fmap (fmap asBool) values)- CTArray ->- D.fromVector (fmap (fmap asArray) values)- CTMixed ->- D.fromVector values--detectColType :: V.Vector (Maybe Value) -> ColType-detectColType vals =- case nonMissing of- [] -> CTMixed- vs- | all isString vs -> CTString- | all isNumber vs -> CTNumber- | all isBool vs -> CTBool- | all isArray vs -> CTArray- | otherwise -> CTMixed- where- nonMissing = catMaybes (V.toList vals)-- isString (String _) = True- isString _ = False-- isNumber (Number _) = True- isNumber _ = False-- isBool (Bool _) = True- isBool _ = False-- isArray (Array _) = True- isArray _ = False--asText :: Value -> Text-asText (String s) = s-asText v = T.pack (show v)--asDouble :: Value -> Double-asDouble (Number s) = toRealFloat @Double s-asDouble v = error $ "asDouble: non-number value: " <> show v--asBool :: Value -> Bool-asBool (Bool b) = b-asBool v = error $ "asBool: non-bool value: " <> show v--asArray :: Value -> V.Vector Value-asArray (Array a) = a-asArray v = error $ "asArray: non-array value: " <> show v
− src/DataFrame/IO/Parquet.hs
@@ -1,735 +0,0 @@-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE MonoLocalBinds #-}-{-# LANGUAGE NumericUnderscores #-}-{-# LANGUAGE OverloadedRecordDot #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}--module DataFrame.IO.Parquet where--import Control.Exception (throw, try)-import Control.Monad-import Control.Monad.IO.Class (MonadIO (..))-import Data.Aeson (FromJSON (..), eitherDecodeStrict, withObject, (.:))-import Data.Bits (Bits (shiftL), (.|.))-import qualified Data.ByteString as BS-import Data.Either (fromRight)-import Data.Functor ((<&>))-import Data.Int (Int32, Int64)-import Data.List (foldl', transpose)-import qualified Data.List as L-import qualified Data.Map as Map-import qualified Data.Text as T-import Data.Text.Encoding (encodeUtf8)-import Data.Time (UTCTime)-import Data.Time.Clock.POSIX (posixSecondsToUTCTime)-import qualified Data.Vector as Vector-import qualified Data.Vector.Unboxed as VU-import DataFrame.Errors (DataFrameException (ColumnsNotFoundException))-import DataFrame.IO.Parquet.Page (- PageDecoder,- UnboxedPageDecoder,- boolDecoder,- byteArrayDecoder,- doubleDecoder,- fixedLenByteArrayDecoder,- floatDecoder,- int32Decoder,- int64Decoder,- int96Decoder,- readPages,- )-import DataFrame.IO.Parquet.Seeking (- FileBufferedOrSeekable,- ForceNonSeekable,- withFileBufferedOrSeekable,- )-import DataFrame.IO.Parquet.Thrift (- ColumnChunk (..),- DecimalType (..),- FileMetadata (..),- LogicalType (..),- RowGroup (..),- ThriftType (..),- TimeUnit (..),- TimestampType (..),- unField,- )-import DataFrame.IO.Parquet.Utils (- ColumnDescription (..),- foldNonNullable,- foldNonNullableUnboxed,- foldNullable,- foldNullableUnboxed,- foldRepeated,- foldRepeatedUnboxed,- generateColumnDescriptions,- getColumnNames,- )-import DataFrame.IO.Utils.RandomAccess (- RandomAccess (..),- ReaderIO (runReaderIO),- )-import DataFrame.Internal.Column (Column, Columnable)-import qualified DataFrame.Internal.Column as DI-import DataFrame.Internal.DataFrame (DataFrame (..))-import DataFrame.Internal.Expression (Expr, getColumns)-import DataFrame.Operations.Merge ()-import qualified DataFrame.Operations.Subset as DS-import Network.HTTP.Simple (- getResponseBody,- getResponseStatusCode,- httpBS,- parseRequest,- setRequestHeader,- )-import qualified Pinch-import qualified Streamly.Data.Stream as Stream-import System.Directory (- doesDirectoryExist,- getHomeDirectory,- getTemporaryDirectory,- )-import System.Environment (lookupEnv)-import System.FilePath ((</>))-import System.FilePath.Glob (compile, glob, match)-import System.IO (IOMode (ReadMode))---- Options -------------------------------------------------------------------{- | Options for reading Parquet data.--These options are applied in this order:--1. predicate filtering-2. column projection-3. row range-4. safe column promotion--Column selection for @selectedColumns@ uses leaf column names only.--}-data ParquetReadOptions = ParquetReadOptions- { selectedColumns :: Maybe [T.Text]- {- ^ Columns to keep in the final dataframe. If set, only these columns are returned.- Predicate-referenced columns are read automatically when needed and projected out after filtering.- -}- , predicate :: Maybe (Expr Bool)- -- ^ Optional row filter expression applied before projection.- , rowRange :: Maybe (Int, Int)- -- ^ Optional row slice @(start, end)@ with start-inclusive/end-exclusive semantics.- , safeColumns :: Bool- -- ^ When True, every column is promoted to OptionalColumn after read, regardless of nullability in the schema.- }- deriving (Show)--{- | Default Parquet read options.--Equivalent to:--@-ParquetReadOptions- { selectedColumns = Nothing- , predicate = Nothing- , rowRange = Nothing- , safeColumns = False- }-@--}-defaultParquetReadOptions :: ParquetReadOptions-defaultParquetReadOptions =- ParquetReadOptions- { selectedColumns = Nothing- , predicate = Nothing- , rowRange = Nothing- , safeColumns = False- }---- Public API ----------------------------------------------------------------{- | Read a parquet file from path and load it into a dataframe.--==== __Example__-@-ghci> D.readParquet ".\/data\/mtcars.parquet"-@--}-readParquet :: FilePath -> IO DataFrame-readParquet = readParquetWithOpts defaultParquetReadOptions--{- | Read a Parquet file using explicit read options.--==== __Example__-@-ghci> D.readParquetWithOpts-ghci| (D.defaultParquetReadOptions{D.selectedColumns = Just ["id"], D.rowRange = Just (0, 10)})-ghci| "./tests/data/alltypes_plain.parquet"-@--When @selectedColumns@ is set and @predicate@ references other columns, those predicate columns-are auto-included for decoding, then projected back to the requested output columns.--}-readParquetWithOpts :: ParquetReadOptions -> FilePath -> IO DataFrame-readParquetWithOpts opts path- | isHFUri path = do- paths <- fetchHFParquetFiles path- let optsNoRange = opts{rowRange = Nothing}- dfs <- mapM (_readParquetWithOpts Nothing optsNoRange) paths- pure (applyRowRange opts (mconcat dfs))- | otherwise = _readParquetWithOpts Nothing opts path---- | Internal entry point used by tests to force non-seekable mode.-_readParquetWithOpts ::- ForceNonSeekable -> ParquetReadOptions -> FilePath -> IO DataFrame-_readParquetWithOpts extraConfig opts path =- withFileBufferedOrSeekable extraConfig path ReadMode $ \file ->- runReaderIO (parseParquetWithOpts opts) file--{- | Read Parquet files from a directory or glob path.--This is equivalent to calling 'readParquetFilesWithOpts' with 'defaultParquetReadOptions'.--}-readParquetFiles :: FilePath -> IO DataFrame-readParquetFiles = readParquetFilesWithOpts defaultParquetReadOptions--{- | Read multiple Parquet files (directory or glob) using explicit options.--If @path@ is a directory, all non-directory entries are read.-If @path@ is a glob, matching files are read.--For multi-file reads, @rowRange@ is applied once after concatenation (global range semantics).--==== __Example__-@-ghci> D.readParquetFilesWithOpts-ghci| (D.defaultParquetReadOptions{D.selectedColumns = Just ["id"], D.rowRange = Just (0, 5)})-ghci| "./tests/data/alltypes_plain*.parquet"-@--}-readParquetFilesWithOpts :: ParquetReadOptions -> FilePath -> IO DataFrame-readParquetFilesWithOpts opts path- | isHFUri path = do- files <- fetchHFParquetFiles path- let optsWithoutRowRange = opts{rowRange = Nothing}- dfs <- mapM (_readParquetWithOpts Nothing optsWithoutRowRange) files- pure (applyRowRange opts (mconcat dfs))- | otherwise = do- isDir <- doesDirectoryExist path-- let pat = if isDir then path </> "*.parquet" else path-- matches <- glob pat-- files <- filterM (fmap not . doesDirectoryExist) matches-- case files of- [] ->- error $- "readParquetFiles: no parquet files found for " ++ path- _ -> do- let optsWithoutRowRange = opts{rowRange = Nothing}- dfs <- mapM (readParquetWithOpts optsWithoutRowRange) files- pure (applyRowRange opts (mconcat dfs))---- Core parsing pipeline -----------------------------------------------------{- | Parse a Parquet file via the 'RandomAccess' handle, applying all-read options. This is the central parsing entry point used by-'_readParquetWithOpts'.--}-parseParquetWithOpts ::- (RandomAccess m, MonadIO m) =>- ParquetReadOptions ->- m DataFrame-parseParquetWithOpts opts = do- metadata <- parseFileMetadata-- let schemaElems = unField metadata.schema- allNames = getColumnNames (drop 1 schemaElems)- leafNames = L.nub (map (last . T.splitOn ".") allNames)- predicateColumns = maybe [] (L.nub . getColumns) (predicate opts)- selectedColumnsForRead = case selectedColumns opts of- Nothing -> Nothing- Just selected -> Just (L.nub (selected ++ predicateColumns))-- -- TODO: When selectedColumnsForRead is Just, pass the set of required- -- column indices into the chunk parsers so that RandomAccess reads are- -- skipped for columns not in the selection, rather than decoding all- -- columns and projecting afterward.-- -- TODO: When rowRange is set, compute cumulative row offsets from- -- rg_num_rows in each RowGroup and skip any group whose row interval does- -- not overlap the requested range, avoiding all decoding for those groups.-- -- TODO: When predicate is set, inspect cmd_statistics min/max values for- -- predicate-referenced columns in each RowGroup and skip groups where- -- statistics prove the predicate cannot be satisfied.-- -- Validate selected columns- case selectedColumnsForRead of- Nothing -> pure ()- Just requested ->- let missing = requested L.\\ leafNames- in unless (L.null missing) $- liftIO $- throw- ( ColumnsNotFoundException- missing- "readParquetWithOpts"- leafNames- )-- let descriptions = generateColumnDescriptions schemaElems- chunks = columnChunksForAll metadata- nCols = length chunks- nDescs = length descriptions-- unless (nCols == nDescs) $- error $- "Column count mismatch: got "- <> show nCols- <> " columns but schema implied "- <> show nDescs- <> " columns"-- -- Some files omit the top-level num_rows field; fall back to summing row-group counts.- let topLevelRows = fromIntegral . unField $ metadata.num_rows :: Int- rgRows =- sum $ map (fromIntegral . unField . rg_num_rows) (unField metadata.row_groups) ::- Int- vectorLength = if topLevelRows > 0 then topLevelRows else rgRows-- rawCols <- zipWithM (parseColumnChunks vectorLength) chunks descriptions-- let finalCols = zipWith applyDescLogicalType descriptions rawCols- indices = Map.fromList $ zip allNames [0 ..]- dimensions = (vectorLength, length finalCols)-- let df =- DataFrame- (Vector.fromListN (length finalCols) finalCols)- indices- dimensions- Map.empty-- return (applyReadOptions opts df)--{- | Parse the file-level Thrift metadata from the Parquet file footer.-Validates the trailing 4-byte magic marker (\"PAR1\") before decoding.--}-parseFileMetadata :: (RandomAccess m) => m FileMetadata-parseFileMetadata = do- footerBytes <- readSuffix 8- let magic = BS.drop 4 footerBytes- when (magic /= "PAR1") $- error- ( "Not a valid Parquet file: expected magic bytes \"PAR1\", got "- ++ show magic- )- let size = getMetadataSize footerBytes- rawMetadata <- readSuffix (size + 8) <&> BS.take size- case Pinch.decode Pinch.compactProtocol rawMetadata of- Left e -> error $ "Failed to parse Parquet metadata: " ++ show e- Right metadata -> return metadata- where- getMetadataSize footer =- let sizes :: [Int]- sizes = map (fromIntegral . BS.index footer) [0 .. 3]- in foldl' (.|.) 0 $ zipWith shiftL sizes [0, 8 .. 24]---- | Read the file metadata from a Parquet file at the given path.-readMetadataFromPath :: FilePath -> IO FileMetadata-readMetadataFromPath path =- withFileBufferedOrSeekable Nothing path ReadMode $- runReaderIO parseFileMetadata---- | Read only the file metadata from an open 'FileBufferedOrSeekable' handle.-readMetadataFromHandle :: FileBufferedOrSeekable -> IO FileMetadata-readMetadataFromHandle = runReaderIO parseFileMetadata---- | Collect column chunks per column (transposed across all row groups).-columnChunksForAll :: FileMetadata -> [[ColumnChunk]]-columnChunksForAll =- transpose . map (unField . rg_columns) . unField . row_groups---- | Dispatch a column's chunks to the correct decoder path.-parseColumnChunks ::- (RandomAccess m, MonadIO m) =>- Int ->- [ColumnChunk] ->- ColumnDescription ->- m Column-parseColumnChunks totalRows chunks description- | description.maxRepetitionLevel == 0 && description.maxDefinitionLevel == 0 =- getNonNullableColumn totalRows description chunks- | description.maxRepetitionLevel == 0 =- getNullableColumn totalRows description chunks- | otherwise =- getRepeatedColumn description chunks---- | Decode a required (non-nullable, non-repeated) column.-getNonNullableColumn ::- forall m.- (RandomAccess m, MonadIO m) =>- Int ->- ColumnDescription ->- [ColumnChunk] ->- m Column-getNonNullableColumn totalRows description chunks =- case description.colElementType of- Just (BOOLEAN _) -> unboxedGo boolDecoder- Just (INT32 _) -> unboxedGo int32Decoder- Just (INT64 _) -> unboxedGo int64Decoder- Just (INT96 _) -> go int96Decoder- Just (FLOAT _) -> unboxedGo floatDecoder- Just (DOUBLE _) -> unboxedGo doubleDecoder- Just (BYTE_ARRAY _) -> go byteArrayDecoder- Just (FIXED_LEN_BYTE_ARRAY _) -> case description.typeLength of- Nothing -> error "FIXED_LEN_BYTE_ARRAY requires type_length to be set"- Just tl -> go (fixedLenByteArrayDecoder (fromIntegral tl))- Nothing -> error "Column has no Parquet type"- where- go ::- forall a.- (Columnable a) =>- PageDecoder a ->- m Column- go decoder =- foldNonNullable totalRows $- (\(vs, _, _) -> vs)- <$> Stream.unfoldEach (readPages description decoder) (Stream.fromList chunks)-- unboxedGo ::- forall a.- (Columnable a, VU.Unbox a) =>- UnboxedPageDecoder a ->- m Column- unboxedGo decoder =- foldNonNullableUnboxed totalRows $- (\(vs, _, _) -> vs)- <$> Stream.unfoldEach- (readPages description decoder)- (Stream.fromList chunks)---- | Decode an optional (nullable) column.-getNullableColumn ::- forall m.- (RandomAccess m, MonadIO m) =>- Int ->- ColumnDescription ->- [ColumnChunk] ->- m Column-getNullableColumn totalRows description chunks =- case description.colElementType of- Just (BOOLEAN _) -> unboxedGo boolDecoder- Just (INT32 _) -> unboxedGo int32Decoder- Just (INT64 _) -> unboxedGo int64Decoder- Just (INT96 _) -> go int96Decoder- Just (FLOAT _) -> unboxedGo floatDecoder- Just (DOUBLE _) -> unboxedGo doubleDecoder- Just (BYTE_ARRAY _) -> go byteArrayDecoder- Just (FIXED_LEN_BYTE_ARRAY _) -> case description.typeLength of- Nothing -> error "FIXED_LEN_BYTE_ARRAY requires type_length to be set"- Just tl -> go (fixedLenByteArrayDecoder (fromIntegral tl))- Nothing -> error "Column has no Parquet type"- where- maxDef :: Int- maxDef = fromIntegral description.maxDefinitionLevel-- go ::- forall a.- (Columnable a) =>- PageDecoder a ->- m Column- go decoder =- foldNullable maxDef totalRows $- (\(vs, ds, _) -> (vs, ds))- <$> Stream.unfoldEach (readPages description decoder) (Stream.fromList chunks)- unboxedGo ::- forall a.- (Columnable a, VU.Unbox a) =>- UnboxedPageDecoder a ->- m Column- unboxedGo decoder =- foldNullableUnboxed maxDef totalRows $- (\(vs, ds, _) -> (vs, ds))- <$> Stream.unfoldEach- (readPages description decoder)- (Stream.fromList chunks)---- | Decode a repeated (list/nested) column.-getRepeatedColumn ::- forall m.- (RandomAccess m, MonadIO m) =>- ColumnDescription ->- [ColumnChunk] ->- m Column-getRepeatedColumn description chunks =- case description.colElementType of- Just (BOOLEAN _) -> unboxedGo boolDecoder- Just (INT32 _) -> unboxedGo int32Decoder- Just (INT64 _) -> unboxedGo int64Decoder- Just (INT96 _) -> go int96Decoder- Just (FLOAT _) -> unboxedGo floatDecoder- Just (DOUBLE _) -> unboxedGo doubleDecoder- Just (BYTE_ARRAY _) -> go byteArrayDecoder- Just (FIXED_LEN_BYTE_ARRAY _) -> case description.typeLength of- Nothing -> error "FIXED_LEN_BYTE_ARRAY requires type_length to be set"- Just tl -> go (fixedLenByteArrayDecoder (fromIntegral tl))- Nothing -> error "Column has no Parquet type"- where- maxRep :: Int- maxRep = fromIntegral description.maxRepetitionLevel- maxDef :: Int- maxDef = fromIntegral description.maxDefinitionLevel-- go ::- forall a.- ( Columnable a- , Columnable (Maybe [Maybe a])- , Columnable (Maybe [Maybe [Maybe a]])- , Columnable (Maybe [Maybe [Maybe [Maybe a]]])- ) =>- PageDecoder a ->- m Column- go decoder =- foldRepeated maxRep maxDef $- Stream.unfoldEach (readPages description decoder) (Stream.fromList chunks)-- unboxedGo ::- forall a.- ( VU.Unbox a- , Columnable a- , Columnable (Maybe [Maybe a])- , Columnable (Maybe [Maybe [Maybe a]])- , Columnable (Maybe [Maybe [Maybe [Maybe a]]])- ) =>- UnboxedPageDecoder a ->- m Column- unboxedGo decoder =- foldRepeatedUnboxed maxRep maxDef $- Stream.unfoldEach- (readPages description decoder)- (Stream.fromList chunks)---- Options application -------------------------------------------------------applyRowRange :: ParquetReadOptions -> DataFrame -> DataFrame-applyRowRange opts df =- maybe df (`DS.range` df) (rowRange opts)--applySelectedColumns :: ParquetReadOptions -> DataFrame -> DataFrame-applySelectedColumns opts df =- maybe df (`DS.select` df) (selectedColumns opts)--applyPredicate :: ParquetReadOptions -> DataFrame -> DataFrame-applyPredicate opts df =- maybe df (`DS.filterWhere` df) (predicate opts)--applySafeRead :: ParquetReadOptions -> DataFrame -> DataFrame-applySafeRead opts df- | safeColumns opts = df{columns = Vector.map DI.ensureOptional (columns df)}- | otherwise = df--applyReadOptions :: ParquetReadOptions -> DataFrame -> DataFrame-applyReadOptions opts =- applySafeRead opts- . applyRowRange opts- . applySelectedColumns opts- . applyPredicate opts---- Logical type conversion ---------------------------------------------------{- | Apply a column-description's logical type annotation to convert raw-decoded values (e.g. millisecond integers → 'UTCTime').--}-applyDescLogicalType :: ColumnDescription -> DI.Column -> DI.Column-applyDescLogicalType desc = applyLogicalType (colLogicalType desc)--applyLogicalType :: Maybe LogicalType -> DI.Column -> DI.Column-applyLogicalType (Just (LT_TIMESTAMP f)) col =- let ts = unField f- unit = unField ts.timestamp_unit- divisor = case unit of- MILLIS _ -> 1_000- MICROS _ -> 1_000_000- NANOS _ -> 1_000_000_000- in fromRight col $- DI.mapColumn- (microsecondsToUTCTime . (* (1_000_000 `div` divisor)))- col-applyLogicalType (Just (LT_DECIMAL f)) col =- let dt = unField f- scale = unField dt.decimal_scale- precision = unField dt.decimal_precision- in if precision <= 9- then case DI.toVector @Int32 @VU.Vector col of- Right xs ->- DI.fromUnboxedVector $- VU.map (\raw -> fromIntegral @Int32 @Double raw / 10 ^ scale) xs- Left _ -> col- else- if precision <= 18- then case DI.toVector @Int64 @VU.Vector col of- Right xs ->- DI.fromUnboxedVector $- VU.map (\raw -> fromIntegral @Int64 @Double raw / 10 ^ scale) xs- Left _ -> col- else col-applyLogicalType _ col = col--microsecondsToUTCTime :: Int64 -> UTCTime-microsecondsToUTCTime us =- posixSecondsToUTCTime (fromIntegral us / 1_000_000)---- HuggingFace support -------------------------------------------------------data HFRef = HFRef- { hfOwner :: T.Text- , hfDataset :: T.Text- , hfGlob :: T.Text- }--data HFParquetFile = HFParquetFile- { hfpUrl :: T.Text- , hfpConfig :: T.Text- , hfpSplit :: T.Text- , hfpFilename :: T.Text- }- deriving (Show)--instance FromJSON HFParquetFile where- parseJSON = withObject "HFParquetFile" $ \o ->- HFParquetFile- <$> o .: "url"- <*> o .: "config"- <*> o .: "split"- <*> o .: "filename"--newtype HFParquetResponse = HFParquetResponse {hfParquetFiles :: [HFParquetFile]}--instance FromJSON HFParquetResponse where- parseJSON = withObject "HFParquetResponse" $ \o ->- HFParquetResponse <$> o .: "parquet_files"--isHFUri :: FilePath -> Bool-isHFUri = L.isPrefixOf "hf://"--parseHFUri :: FilePath -> Either String HFRef-parseHFUri path =- let stripped = drop (length ("hf://datasets/" :: String)) path- in case T.splitOn "/" (T.pack stripped) of- (owner : dataset : rest)- | not (null rest) ->- Right $ HFRef owner dataset (T.intercalate "/" rest)- _ ->- Left $ "Invalid hf:// URI (expected hf://datasets/owner/dataset/glob): " ++ path--getHFToken :: IO (Maybe BS.ByteString)-getHFToken = do- envToken <- lookupEnv "HF_TOKEN"- case envToken of- Just t -> pure (Just (encodeUtf8 (T.pack t)))- Nothing -> do- home <- getHomeDirectory- let tokenPath = home </> ".cache" </> "huggingface" </> "token"- result <- try (BS.readFile tokenPath) :: IO (Either IOError BS.ByteString)- case result of- Right bs -> pure (Just (BS.takeWhile (/= 10) bs))- Left _ -> pure Nothing--{- | Extract the repo-relative path from a HuggingFace download URL.-URL format: https://huggingface.co/datasets/{owner}/{dataset}/resolve/{ref}/{path}-Returns the {path} portion (e.g. "data/train-00000-of-00001.parquet").--}-hfUrlRepoPath :: HFParquetFile -> String-hfUrlRepoPath f =- case T.breakOn "/resolve/" (hfpUrl f) of- (_, rest)- | not (T.null rest) ->- -- Drop "/resolve/", then drop the ref component (up to and including "/")- T.unpack $ T.drop 1 $ T.dropWhile (/= '/') $ T.drop (T.length "/resolve/") rest- _ ->- T.unpack (hfpConfig f) </> T.unpack (hfpSplit f) </> T.unpack (hfpFilename f)--matchesGlob :: T.Text -> HFParquetFile -> Bool-matchesGlob g f = match (compile (T.unpack g)) (hfUrlRepoPath f)--resolveHFUrls :: Maybe BS.ByteString -> HFRef -> IO [HFParquetFile]-resolveHFUrls mToken ref = do- let dataset = hfOwner ref <> "/" <> hfDataset ref- let apiUrl = "https://datasets-server.huggingface.co/parquet?dataset=" ++ T.unpack dataset- req0 <- parseRequest apiUrl- let req = case mToken of- Nothing -> req0- Just tok -> setRequestHeader "Authorization" ["Bearer " <> tok] req0- resp <- httpBS req- let status = getResponseStatusCode resp- when (status /= 200) $- ioError $- userError $- "HuggingFace API returned status "- ++ show status- ++ " for dataset "- ++ T.unpack dataset- case eitherDecodeStrict (getResponseBody resp) of- Left err -> ioError $ userError $ "Failed to parse HF API response: " ++ err- Right hfResp -> pure $ filter (matchesGlob (hfGlob ref)) (hfParquetFiles hfResp)--downloadHFFiles :: Maybe BS.ByteString -> [HFParquetFile] -> IO [FilePath]-downloadHFFiles mToken files = do- tmpDir <- getTemporaryDirectory- forM files $ \f -> do- -- Derive a collision-resistant temp name from the URL path components- let fname = case (hfpConfig f, hfpSplit f) of- (c, s) | T.null c && T.null s -> T.unpack (hfpFilename f)- (c, s) -> T.unpack c <> "_" <> T.unpack s <> "_" <> T.unpack (hfpFilename f)- let destPath = tmpDir </> fname- req0 <- parseRequest (T.unpack (hfpUrl f))- let req = case mToken of- Nothing -> req0- Just tok -> setRequestHeader "Authorization" ["Bearer " <> tok] req0- resp <- httpBS req- let status = getResponseStatusCode resp- when (status /= 200) $- ioError $- userError $- "Failed to download " ++ T.unpack (hfpUrl f) ++ " (HTTP " ++ show status ++ ")"- BS.writeFile destPath (getResponseBody resp)- pure destPath---- | True when the path contains glob wildcard characters.-hasGlob :: T.Text -> Bool-hasGlob = T.any (\c -> c == '*' || c == '?' || c == '[')--{- | Build the direct HF repo download URL for a path with no wildcards.-Format: https://huggingface.co/datasets/{owner}/{dataset}/resolve/main/{path}--}-directHFUrl :: HFRef -> T.Text-directHFUrl ref =- "https://huggingface.co/datasets/"- <> hfOwner ref- <> "/"- <> hfDataset ref- <> "/resolve/main/"- <> hfGlob ref--fetchHFParquetFiles :: FilePath -> IO [FilePath]-fetchHFParquetFiles uri = do- ref <- case parseHFUri uri of- Left err -> ioError (userError err)- Right r -> pure r- mToken <- getHFToken- if hasGlob (hfGlob ref)- then do- hfFiles <- resolveHFUrls mToken ref- when (null hfFiles) $- ioError $- userError $- "No parquet files found for " ++ uri- downloadHFFiles mToken hfFiles- else do- -- Direct repo file download — no datasets-server needed- let url = directHFUrl ref- let filename = last $ T.splitOn "/" (hfGlob ref)- downloadHFFiles mToken [HFParquetFile url "" "" filename]
− src/DataFrame/IO/Parquet/Binary.hs
@@ -1,141 +0,0 @@-{-# LANGUAGE TypeApplications #-}--module DataFrame.IO.Parquet.Binary where--import Control.Exception (bracketOnError)-import Control.Monad-import Data.Bits-import qualified Data.ByteString as BS-import qualified Data.ByteString.Unsafe as BSU-import Data.Char-import Data.IORef-import Data.Int-import Data.Word-import qualified Foreign.Marshal.Alloc as Foreign-import qualified Foreign.Ptr as Foreign-import qualified Foreign.Storable as Foreign--readUVarInt :: BS.ByteString -> (Word64, BS.ByteString)-readUVarInt xs = loop xs 0 0 0- where- {-- Each input byte contributes:- - lower 7 payload bits- - The high bit (0x80) is the continuation flag: 1 = more bytes follow, 0 = last byte- Why the magic number 10: For a 64‑bit integer we need at most ceil(64 / 7) = 10 bytes- -}- loop :: BS.ByteString -> Word64 -> Int -> Int -> (Word64, BS.ByteString)- loop bs result _ 10 = (result, bs)- loop xs' result shiftAmt i = case BS.uncons xs' of- Nothing -> error "readUVarInt: not enough input bytes"- Just (b, bs') ->- if b < 0x80- then (result .|. (fromIntegral b `shiftL` shiftAmt), bs')- else- let payloadBits = fromIntegral (b .&. 0x7f) :: Word64- in loop bs' (result .|. (payloadBits `shiftL` shiftAmt)) (shiftAmt + 7) (i + 1)--readVarIntFromBytes :: (Integral a) => BS.ByteString -> (a, BS.ByteString)-readVarIntFromBytes bs = (fromIntegral n, remainder)- where- (n, remainder) = loop 0 0 bs- loop shiftAmt result bs' = case BS.uncons bs' of- Nothing -> (result, BS.empty)- Just (x, xs) ->- let res = result .|. (fromIntegral (x .&. 0x7f) :: Integer) `shiftL` shiftAmt- in if x .&. 0x80 /= 0x80 then (res, xs) else loop (shiftAmt + 7) res xs--readIntFromBytes :: (Integral a) => BS.ByteString -> (a, BS.ByteString)-readIntFromBytes bs =- let (n, remainder) = readVarIntFromBytes bs- u = fromIntegral n :: Word32- in ( fromIntegral $ (fromIntegral (u `shiftR` 1) :: Int32) .^. (-(n .&. 1))- , remainder- )--readInt32FromBytes :: BS.ByteString -> (Int32, BS.ByteString)-readInt32FromBytes bs =- let (n', remainder) = readVarIntFromBytes @Int64 bs- n = fromIntegral n' :: Int32- u = fromIntegral n :: Word32- in ((fromIntegral (u `shiftR` 1) :: Int32) .^. (-(n .&. 1)), remainder)--readAndAdvance :: IORef Int -> BS.ByteString -> IO Word8-readAndAdvance bufferPos buffer = do- pos <- readIORef bufferPos- let b = BS.index buffer pos- modifyIORef' bufferPos (+ 1)- return b--readVarIntFromBuffer :: (Integral a) => BS.ByteString -> IORef Int -> IO a-readVarIntFromBuffer buf bufferPos = do- start <- readIORef bufferPos- let loop i shiftAmt result = do- b <- readAndAdvance bufferPos buf- let res = result .|. (fromIntegral (b .&. 0x7f) :: Integer) `shiftL` shiftAmt- if b .&. 0x80 /= 0x80- then return res- else loop (i + 1) (shiftAmt + 7) res- fromIntegral <$> loop start 0 0--readIntFromBuffer :: (Integral a) => BS.ByteString -> IORef Int -> IO a-readIntFromBuffer buf bufferPos = do- n <- readVarIntFromBuffer buf bufferPos- let u = fromIntegral n :: Word32- return $ fromIntegral $ (fromIntegral (u `shiftR` 1) :: Int32) .^. (-(n .&. 1))--readInt32FromBuffer :: BS.ByteString -> IORef Int -> IO Int32-readInt32FromBuffer buf bufferPos = do- n <- (fromIntegral <$> readVarIntFromBuffer @Int64 buf bufferPos) :: IO Int32- let u = fromIntegral n :: Word32- return $ (fromIntegral (u `shiftR` 1) :: Int32) .^. (-(n .&. 1))--readString :: BS.ByteString -> IORef Int -> IO String-readString buf pos = do- nameSize <- readVarIntFromBuffer @Int buf pos- replicateM nameSize (chr . fromIntegral <$> readAndAdvance pos buf)--readByteStringFromBytes :: BS.ByteString -> (BS.ByteString, BS.ByteString)-readByteStringFromBytes xs =- let- (size, remainder) = readVarIntFromBytes @Int xs- in- BS.splitAt size remainder--readByteString :: BS.ByteString -> IORef Int -> IO BS.ByteString-readByteString buf pos = do- size <- readVarIntFromBuffer @Int buf pos- fillByteStringByWord8 size (\_ -> readAndAdvance pos buf)--readByteString' :: BS.ByteString -> Int64 -> IO BS.ByteString-readByteString' buf size =- fillByteStringByWord8- (fromIntegral size)- ((`readSingleByte` buf) . fromIntegral)--{- | Allocate a fixed-size buffer, repeat the action on each index.-Fill it into the buffer to get a ByteString.--}-fillByteStringByWord8 :: Int -> (Int -> IO Word8) -> IO BS.ByteString-fillByteStringByWord8 size getByte = do- bracketOnError- (Foreign.mallocBytes size :: IO (Foreign.Ptr Word8))- Foreign.free- -- \^ ensures p is freed if (IO Word8) throws.- ( \p -> do- fill 0 p- BSU.unsafePackCStringFinalizer p size (Foreign.free p)- )- where- fill i p- | i >= size = pure ()- | otherwise = getByte i >>= Foreign.pokeByteOff p i >> fill (i + 1) p-{-# INLINE fillByteStringByWord8 #-}--readSingleByte :: Int64 -> BS.ByteString -> IO Word8-readSingleByte pos buffer = return $ BS.index buffer (fromIntegral pos)--readNoAdvance :: IORef Int -> BS.ByteString -> IO Word8-readNoAdvance bufferPos buffer = do- pos <- readIORef bufferPos- return $ BS.index buffer pos
− src/DataFrame/IO/Parquet/Decompress.hs
@@ -1,32 +0,0 @@-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
@@ -1,164 +0,0 @@-{-# LANGUAGE BangPatterns #-}--module DataFrame.IO.Parquet.Dictionary (DictVals (..), readDictVals, decodeRLEBitPackedHybrid) where--import Data.Bits-import qualified Data.ByteString as BS-import qualified Data.ByteString.Unsafe as BSU-import Data.Int (Int32, Int64)-import qualified Data.Text as T-import Data.Text.Encoding-import Data.Time (UTCTime)-import qualified Data.Vector as V-import Data.Word-import DataFrame.IO.Parquet.Binary (readUVarInt)-import DataFrame.IO.Parquet.Thrift (ThriftType (..))-import DataFrame.IO.Parquet.Time (int96ToUTCTime)-import DataFrame.Internal.Binary (- littleEndianInt32,- littleEndianWord32,- littleEndianWord64,- )-import GHC.Float--data DictVals- = DBool (V.Vector Bool)- | DInt32 (V.Vector Int32)- | DInt64 (V.Vector Int64)- | DInt96 (V.Vector UTCTime)- | DFloat (V.Vector Float)- | DDouble (V.Vector Double)- | DText (V.Vector T.Text)- deriving (Show, Eq)--{- | Decode the values from a dictionary page.--The @numVals@ argument is the entry count declared in the dictionary page-header. It is used to limit BOOLEAN decoding (1-bit-per-value encoding has-no natural delimiter).--The @typeLength@ argument is only meaningful for FIXED_LEN_BYTE_ARRAY: it is-the byte-width of each individual dictionary entry, NOT the total number of-entries. Passing @numVals@ here (the old behaviour) would cause it to be-misread as an element size, yielding a dictionary that is far too small.--}-readDictVals :: ThriftType -> BS.ByteString -> Int32 -> Maybe Int32 -> DictVals-readDictVals (BOOLEAN _) bs count _ = DBool (V.fromList (take (fromIntegral count) $ readPageBool bs))-readDictVals (INT32 _) bs _ _ = DInt32 (V.fromList (readPageInt32 bs))-readDictVals (INT64 _) bs _ _ = DInt64 (V.fromList (readPageInt64 bs))-readDictVals (INT96 _) bs _ _ = DInt96 (V.fromList (readPageInt96Times bs))-readDictVals (FLOAT _) bs _ _ = DFloat (V.fromList (readPageFloat bs))-readDictVals (DOUBLE _) bs _ _ = DDouble (V.fromList (readPageWord64 bs))-readDictVals (BYTE_ARRAY _) bs _ _ = DText (V.fromList (readPageBytes bs))-readDictVals (FIXED_LEN_BYTE_ARRAY _) bs _ (Just len) =- DText (V.fromList (readPageFixedBytes bs (fromIntegral len)))-readDictVals t _ _ _ = error $ "Unsupported dictionary type: " ++ show t--readPageInt32 :: BS.ByteString -> [Int32]-readPageInt32 xs- | BS.null xs = []- | otherwise = littleEndianInt32 (BS.take 4 xs) : readPageInt32 (BS.drop 4 xs)--readPageWord64 :: BS.ByteString -> [Double]-readPageWord64 xs- | BS.null xs = []- | otherwise =- castWord64ToDouble (littleEndianWord64 (BS.take 8 xs))- : readPageWord64 (BS.drop 8 xs)--readPageBytes :: BS.ByteString -> [T.Text]-readPageBytes xs- | BS.null xs = []- | otherwise =- let lenBytes = fromIntegral (littleEndianInt32 $ BS.take 4 xs)- totalBytesRead = lenBytes + 4- in decodeUtf8Lenient (BS.take lenBytes (BS.drop 4 xs))- : readPageBytes (BS.drop totalBytesRead xs)--readPageBool :: BS.ByteString -> [Bool]-readPageBool bs =- concatMap (\b -> map (\i -> (b `shiftR` i) .&. 1 == 1) [0 .. 7]) (BS.unpack bs)--readPageInt64 :: BS.ByteString -> [Int64]-readPageInt64 xs- | BS.null xs = []- | otherwise =- fromIntegral (littleEndianWord64 (BS.take 8 xs)) : readPageInt64 (BS.drop 8 xs)--readPageFloat :: BS.ByteString -> [Float]-readPageFloat xs- | BS.null xs = []- | otherwise =- castWord32ToFloat (littleEndianWord32 (BS.take 4 xs))- : readPageFloat (BS.drop 4 xs)--readNInt96Times :: Int -> BS.ByteString -> ([UTCTime], BS.ByteString)-readNInt96Times 0 bs = ([], bs)-readNInt96Times k bs =- let timestamp96 = BS.take 12 bs- utcTime = int96ToUTCTime timestamp96- bs' = BS.drop 12 bs- (times, rest) = readNInt96Times (k - 1) bs'- in (utcTime : times, rest)--readPageInt96Times :: BS.ByteString -> [UTCTime]-readPageInt96Times bs- | BS.null bs = []- | otherwise =- let (times, _) = readNInt96Times (BS.length bs `div` 12) bs- in times--readPageFixedBytes :: BS.ByteString -> Int -> [T.Text]-readPageFixedBytes xs len- | BS.null xs = []- | otherwise =- decodeUtf8Lenient (BS.take len xs) : readPageFixedBytes (BS.drop len xs) len--unpackBitPacked :: Int -> Int -> BS.ByteString -> ([Word32], BS.ByteString)-unpackBitPacked bw count bs- | count <= 0 = ([], bs)- | BS.null bs = ([], bs)- | otherwise =- let totalBytes = (bw * count + 7) `div` 8- chunk = BS.take totalBytes bs- rest = BS.drop totalBytes bs- in (extractBits bw count chunk, rest)---- | LSB-first bit accumulator: reads each byte once with no intermediate ByteString allocation.-extractBits :: Int -> Int -> BS.ByteString -> [Word32]-extractBits bw count bs = go 0 (0 :: Word64) 0 count- where- !mask = if bw == 32 then maxBound else (1 `shiftL` bw) - 1 :: Word64- !len = BS.length bs- go !byteIdx !acc !accBits !remaining- | remaining <= 0 = []- | accBits >= bw =- fromIntegral (acc .&. mask)- : go byteIdx (acc `shiftR` bw) (accBits - bw) (remaining - 1)- | byteIdx >= len = []- | otherwise =- let b = fromIntegral (BSU.unsafeIndex bs byteIdx) :: Word64- in go (byteIdx + 1) (acc .|. (b `shiftL` accBits)) (accBits + 8) remaining--decodeRLEBitPackedHybrid :: Int -> BS.ByteString -> ([Word32], BS.ByteString)-decodeRLEBitPackedHybrid bitWidth bs- | bitWidth == 0 = ([0], bs)- | BS.null bs = ([], bs)- | otherwise =- -- readUVarInt is evaluated here, inside the guard that has already- -- confirmed bs is non-empty. Keeping it in a where clause would cause- -- it to be forced before the BS.null guard under {-# LANGUAGE Strict #-}.- let (hdr64, afterHdr) = readUVarInt bs- isPacked = (hdr64 .&. 1) == 1- in if isPacked- then- let groups = fromIntegral (hdr64 `shiftR` 1) :: Int- totalVals = groups * 8- in unpackBitPacked bitWidth totalVals afterHdr- else- let mask = if bitWidth == 32 then maxBound else (1 `shiftL` bitWidth) - 1- runLen = fromIntegral (hdr64 `shiftR` 1) :: Int- nBytes = (bitWidth + 7) `div` 8 :: Int- word32 = littleEndianWord32 (BS.take 4 afterHdr)- value = word32 .&. mask- in (replicate runLen value, BS.drop nBytes afterHdr)
− src/DataFrame/IO/Parquet/Encoding.hs
@@ -1,135 +0,0 @@-{-# LANGUAGE BangPatterns #-}-{-# LANGUAGE CPP #-}--module DataFrame.IO.Parquet.Encoding (- -- Kept from the original Encoding module (used by Levels)- ceilLog2,- bitWidthForMaxLevel,- -- Vector-based RLE/bit-packed decoder (from new parser)- decodeRLEBitPackedHybrid,- extractBitsInto,- fillRun,- decodeDictIndices,-) where--import Control.Monad.ST (ST, runST)-import Data.Bits-import qualified Data.ByteString as BS-import qualified Data.ByteString.Unsafe as BSU-#if !MIN_VERSION_base(4,20,0)-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- | otherwise = 1 + ceilLog2 ((x + 1) `div` 2)--bitWidthForMaxLevel :: Int -> Int-bitWidthForMaxLevel maxLevel = ceilLog2 (maxLevel + 1)---- ------------------------------------------------------------------------------ Vector-based RLE / bit-packed hybrid decoder--- -----------------------------------------------------------------------------decodeRLEBitPackedHybrid ::- -- | Bit width per value (0 = all zeros, use 'VU.replicate')- Int ->- -- | Exact number of values to decode- Int ->- BS.ByteString ->- (VU.Vector Word32, BS.ByteString)-decodeRLEBitPackedHybrid bw need bs- | bw == 0 = (VU.replicate need 0, bs)- | otherwise = runST $ do- mv <- VUM.new need- rest <- go mv 0 bs- dat <- VU.unsafeFreeze mv- return (dat, rest)- where- !mask = if bw == 32 then maxBound else (1 `shiftL` bw) - 1 :: Word32- go :: VUM.STVector s Word32 -> Int -> BS.ByteString -> ST s BS.ByteString- go mv !filled !buf- | filled >= need = return buf- | BS.null buf = return buf- | otherwise =- let (hdr64, afterHdr) = readUVarInt buf- isPacked = (hdr64 .&. 1) == 1- in if isPacked- then do- let groups = fromIntegral (hdr64 `shiftR` 1) :: Int- totalVals = groups * 8- takeN = min (need - filled) totalVals- -- Consume all the bytes for this group even if we- -- only need a subset of the values.- bytesN = (bw * totalVals + 7) `div` 8- (chunk, rest) = BS.splitAt bytesN afterHdr- extractBitsInto bw takeN chunk mv filled- go mv (filled + takeN) rest- else do- let runLen = fromIntegral (hdr64 `shiftR` 1) :: Int- nbytes = (bw + 7) `div` 8- val = littleEndianWord32 (BS.take 4 afterHdr) .&. mask- takeN = min (need - filled) runLen- -- Fill the run directly — no list, no reverse.- fillRun mv filled (filled + takeN) val- go mv (filled + takeN) (BS.drop nbytes afterHdr)-{-# INLINE decodeRLEBitPackedHybrid #-}---- | Fill @mv[start..end-1]@ with @val@ using a bulk @memset@-style write.-fillRun :: VUM.STVector s Word32 -> Int -> Int -> Word32 -> ST s ()-fillRun mv i end = VUM.set (VUM.unsafeSlice i (end - i) mv)-{-# INLINE fillRun #-}--{- | Write @count@ bit-width-@bw@ values from @bs@ into @mv@ starting at-@offset@, reading the byte buffer with a single-pass LSB-first accumulator.-No intermediate list or ByteString allocation.--}-extractBitsInto ::- -- | Bit width- Int ->- -- | Number of values to extract- Int ->- BS.ByteString ->- VUM.STVector s Word32 ->- -- | Write offset into @mv@- Int ->- ST s ()-extractBitsInto bw count bs mv off = go 0 (0 :: Word64) 0 0- where- !mask = if bw == 32 then maxBound else (1 `unsafeShiftL` bw) - 1 :: Word64- !len = BS.length bs- go !byteIdx !acc !accBits !done- | done >= count = return ()- | accBits >= bw = do- VUM.unsafeWrite mv (off + done) (fromIntegral (acc .&. mask))- go byteIdx (acc `unsafeShiftR` bw) (accBits - bw) (done + 1)- | byteIdx >= len = return ()- | otherwise =- let b = fromIntegral (BSU.unsafeIndex bs byteIdx) :: Word64- in go (byteIdx + 1) (acc .|. (b `unsafeShiftL` accBits)) (accBits + 8) done-{-# INLINE extractBitsInto #-}--{- | Decode @need@ dictionary indices from a DATA_PAGE bit-width-prefixed-stream (the first byte encodes the bit-width of all subsequent RLE\/bitpacked-values).--Returns the index vector (as 'Int') and the unconsumed bytes.--}-decodeDictIndices :: Int -> BS.ByteString -> (VU.Vector Int, BS.ByteString)-decodeDictIndices need bs = case BS.uncons bs of- Nothing -> error "decodeDictIndices: empty stream"- Just (w0, rest0) ->- let bw = fromIntegral w0 :: Int- (raw, rest1) = decodeRLEBitPackedHybrid bw need rest0- in (VU.map fromIntegral raw, rest1)-{-# INLINE decodeDictIndices #-}
− src/DataFrame/IO/Parquet/Levels.hs
@@ -1,210 +0,0 @@-{-# LANGUAGE BangPatterns #-}--module DataFrame.IO.Parquet.Levels (- -- Level readers- readLevelsV1,- readLevelsV2,- -- Stitch functions- stitchList,- stitchList2,- stitchList3,-) where--import Control.Monad.ST (runST)-import qualified Data.ByteString as BS-import Data.Int (Int32)-import qualified Data.Vector as VB-import qualified Data.Vector.Unboxed as VU-import qualified Data.Vector.Unboxed.Mutable as VUM-import Data.Word (Word32)-import DataFrame.IO.Parquet.Encoding (- bitWidthForMaxLevel,- decodeRLEBitPackedHybrid,- )-import DataFrame.Internal.Binary (littleEndianWord32)---- ------------------------------------------------------------------------------ Level readers--- -----------------------------------------------------------------------------{- | Convert a 'Word32' level vector to an 'Int' level vector while counting-how many entries equal @maxDef@. Single pass; allocates a single-'VU.Vector Int' of length @VU.length raw@.--}-convertAndCount :: Int -> VU.Vector Word32 -> (VU.Vector Int, Int)-convertAndCount maxDef raw = runST $ do- let !n = VU.length raw- mv <- VUM.unsafeNew n- let !maxDefW = fromIntegral maxDef :: Word32- go !i !nPresent- | i >= n = pure nPresent- | otherwise = do- let !w = VU.unsafeIndex raw i- !d = fromIntegral w :: Int- VUM.unsafeWrite mv i d- if w == maxDefW- then go (i + 1) (nPresent + 1)- else go (i + 1) nPresent- !nPresent <- go 0 0- !out <- VU.unsafeFreeze mv- pure (out, nPresent)--readLevelsV1 ::- -- | Total number of values in the page- Int ->- -- | maxDefinitionLevel- Int ->- -- | maxRepetitionLevel- Int ->- BS.ByteString ->- (VU.Vector Int, VU.Vector Int, Int, BS.ByteString)-readLevelsV1 n maxDef maxRep bs =- let bwRep = bitWidthForMaxLevel maxRep- bwDef = bitWidthForMaxLevel maxDef- (repVec, _, afterRep) = decodeLevelBlock bwRep n bs- (defVec, nPresent, afterDef) = decodeLevelBlock bwDef n afterRep- in (defVec, repVec, nPresent, afterDef)- where- -- For rep block we don't need nPresent; we still get one cheaply.- decodeLevelBlock 0 n' buf = (VU.replicate n' 0, n' * fromEnum (maxDef == 0), buf)- decodeLevelBlock bw n' buf =- let blockLen = fromIntegral (littleEndianWord32 (BS.take 4 buf)) :: Int- blockData = BS.take blockLen (BS.drop 4 buf)- after = BS.drop (4 + blockLen) buf- (raw, _) = decodeRLEBitPackedHybrid bw n' blockData- (out, np) = convertAndCount maxDef raw- in (out, np, after)--readLevelsV2 ::- -- | Total number of values- Int ->- -- | maxDefinitionLevel- Int ->- -- | maxRepetitionLevel- Int ->- -- | Repetition-level byte length (from page header)- Int32 ->- -- | Definition-level byte length (from page header)- Int32 ->- BS.ByteString ->- (VU.Vector Int, VU.Vector Int, Int, BS.ByteString)-readLevelsV2 n maxDef maxRep repLen defLen bs =- let (repBytes, afterRepBytes) = BS.splitAt (fromIntegral repLen) bs- (defBytes, afterDefBytes) = BS.splitAt (fromIntegral defLen) afterRepBytes- bwRep = bitWidthForMaxLevel maxRep- bwDef = bitWidthForMaxLevel maxDef- repVec- | bwRep == 0 = VU.replicate n 0- | otherwise =- let (raw, _) = decodeRLEBitPackedHybrid bwRep n repBytes- (out, _) = convertAndCount maxDef raw- in out- (defVec, nPresent)- | bwDef == 0 = (VU.replicate n 0, n * fromEnum (maxDef == 0))- | otherwise =- let (raw, _) = decodeRLEBitPackedHybrid bwDef n defBytes- in convertAndCount maxDef raw- in (defVec, repVec, nPresent, afterDefBytes)--{- | Stitch a singly-nested list column (@maxRep == 1@) from vector-format-definition and repetition levels plus a compact present-values vector.-Returns one @Maybe [Maybe a]@ per top-level row.--}-stitchList ::- Int ->- VU.Vector Int ->- VU.Vector Int ->- VB.Vector a ->- [Maybe [Maybe a]]-stitchList maxDef repVec defVec values =- map toRow (splitAtRepBound 0 (pairWithValsV maxDef repVec defVec values))- where- toRow [] = Nothing- toRow ((_, d, _) : _) | d == 0 = Nothing- toRow grp = Just [v | (_, _, v) <- grp]--{- | Stitch a doubly-nested list column (@maxRep == 2@).-@defT1@ is the def threshold at which the depth-1 element is present.--}-stitchList2 ::- Int ->- Int ->- VU.Vector Int ->- VU.Vector Int ->- VB.Vector a ->- [Maybe [Maybe [Maybe a]]]-stitchList2 defT1 maxDef repVec defVec values =- map toRow (splitAtRepBound 0 triplets)- where- triplets = pairWithValsV maxDef repVec defVec values- toRow [] = Nothing- toRow ((_, d, _) : _) | d == 0 = Nothing- toRow row = Just (map toOuter (splitAtRepBound 1 row))- toOuter [] = Nothing- toOuter ((_, d, _) : _) | d < defT1 = Nothing- toOuter outer = Just (map toLeaf (splitAtRepBound 2 outer))- toLeaf [] = Nothing- toLeaf ((_, _, v) : _) = v--{- | Stitch a triply-nested list column (@maxRep == 3@).-@defT1@ and @defT2@ are the def thresholds for depth-1 and depth-2-elements respectively.--}-stitchList3 ::- Int ->- Int ->- Int ->- VU.Vector Int ->- VU.Vector Int ->- VB.Vector a ->- [Maybe [Maybe [Maybe [Maybe a]]]]-stitchList3 defT1 defT2 maxDef repVec defVec values =- map toRow (splitAtRepBound 0 triplets)- where- triplets = pairWithValsV maxDef repVec defVec values- toRow [] = Nothing- toRow ((_, d, _) : _) | d == 0 = Nothing- toRow row = Just (map toOuter (splitAtRepBound 1 row))- toOuter [] = Nothing- toOuter ((_, d, _) : _) | d < defT1 = Nothing- toOuter outer = Just (map toMiddle (splitAtRepBound 2 outer))- toMiddle [] = Nothing- toMiddle ((_, d, _) : _) | d < defT2 = Nothing- toMiddle middle = Just (map toLeaf (splitAtRepBound 3 middle))- toLeaf [] = Nothing- toLeaf ((_, _, v) : _) = v---- ------------------------------------------------------------------------------ Internal helpers--- -----------------------------------------------------------------------------{- | Zip rep and def level vectors with a present-values vector, tagging each-position as @Just value@ (when @def == maxDef@) or @Nothing@.-Returns a flat list of @(rep, def, Maybe a)@ triplets for row-splitting.--}-pairWithValsV ::- Int ->- VU.Vector Int ->- VU.Vector Int ->- VB.Vector a ->- [(Int, Int, Maybe a)]-pairWithValsV maxDef repVec defVec values = go 0 0- where- n = VU.length defVec- go i j- | i >= n = []- | otherwise =- let r = VU.unsafeIndex repVec i- d = VU.unsafeIndex defVec i- in if d == maxDef- then (r, d, Just (VB.unsafeIndex values j)) : go (i + 1) (j + 1)- else (r, d, Nothing) : go (i + 1) j--{- | Group a flat triplet list into rows.-A new group begins whenever @rep <= bound@.--}-splitAtRepBound :: Int -> [(Int, Int, Maybe a)] -> [[(Int, Int, Maybe a)]]-splitAtRepBound _ [] = []-splitAtRepBound bound (t : ts) =- let (rest, remaining) = span (\(r, _, _) -> r > bound) ts- in (t : rest) : splitAtRepBound bound remaining
− src/DataFrame/IO/Parquet/Page.hs
@@ -1,352 +0,0 @@-{-# LANGUAGE OverloadedRecordDot #-}-{-# LANGUAGE ScopedTypeVariables #-}--module DataFrame.IO.Parquet.Page (- -- Types- PageDecoder,- UnboxedPageDecoder,- -- Per-type decoders- boolDecoder,- int32Decoder,- int64Decoder,- int96Decoder,- floatDecoder,- doubleDecoder,- byteArrayDecoder,- fixedLenByteArrayDecoder,- -- Page iteration- readPages,-) where--import Control.Monad.IO.Class (MonadIO (liftIO))-import Data.Bits (shiftR, (.&.))-import qualified Data.ByteString as BS-import Data.Int (Int32, Int64)-import Data.Maybe (fromJust, fromMaybe)-import qualified Data.Text as T-import Data.Text.Encoding (decodeUtf8Lenient)-import Data.Time (UTCTime)-import qualified Data.Vector as VB-import qualified Data.Vector.Generic as VG-import qualified Data.Vector.Unboxed as VU-import DataFrame.IO.Parquet.Decompress (decompressData)-import DataFrame.IO.Parquet.Dictionary (- DictVals (..),- readDictVals,- )-import DataFrame.IO.Parquet.Encoding (decodeDictIndices)-import DataFrame.IO.Parquet.Levels (readLevelsV1, readLevelsV2)-import DataFrame.IO.Parquet.Thrift (- ColumnChunk (..),- ColumnMetaData (..),- CompressionCodec,- DataPageHeader (..),- DataPageHeaderV2 (..),- DictionaryPageHeader (..),- Encoding (..),- PageHeader (..),- PageType (..),- ThriftType (..),- unField,- )-import DataFrame.IO.Parquet.Time (int96ToUTCTime)-import DataFrame.IO.Parquet.Utils (ColumnDescription (..))-import DataFrame.IO.Utils.RandomAccess (RandomAccess (..), Range (Range))-import DataFrame.Internal.Binary (- littleEndianInt32,- littleEndianWord32,- littleEndianWord64,- )-import GHC.Float (castWord32ToFloat, castWord64ToDouble)-import Pinch (decodeWithLeftovers)-import qualified Pinch-import Streamly.Internal.Data.Unfold (Step (..), Unfold, mkUnfoldM)---- ------------------------------------------------------------------------------ Types--- -----------------------------------------------------------------------------{- | A type-specific page decoder.-Given the optional dictionary, the page encoding, the number of present-values, and the decompressed value bytes, returns exactly @nPresent@ values.--}-type PageDecoder a =- Maybe DictVals -> Encoding -> Int -> BS.ByteString -> VB.Vector a--type UnboxedPageDecoder a =- Maybe DictVals -> Encoding -> Int -> BS.ByteString -> VU.Vector a---- ------------------------------------------------------------------------------ Per-type decoders--- -----------------------------------------------------------------------------boolDecoder :: UnboxedPageDecoder Bool-boolDecoder mDict enc nPresent bs = case enc of- PLAIN _ -> VU.fromList (readNBool nPresent bs)- RLE_DICTIONARY _ -> unboxedLookupDict mDict nPresent bs getBool- PLAIN_DICTIONARY _ -> unboxedLookupDict mDict nPresent bs getBool- _ -> error ("boolDecoder: unsupported encoding " ++ show enc)- where- getBool (DBool ds) i = ds VB.! i- getBool d _ = error ("boolDecoder: wrong dict type, got " ++ show d)--int32Decoder :: UnboxedPageDecoder Int32-int32Decoder mDict enc nPresent bs = case enc of- PLAIN _ -> VU.convert (readNInt32 nPresent bs)- RLE_DICTIONARY _ -> unboxedLookupDict mDict nPresent bs getInt32- PLAIN_DICTIONARY _ -> unboxedLookupDict mDict nPresent bs getInt32- _ -> error ("int32Decoder: unsupported encoding " ++ show enc)- where- getInt32 (DInt32 ds) i = ds VB.! i- getInt32 d _ = error ("int32Decoder: wrong dict type, got " ++ show d)--int64Decoder :: UnboxedPageDecoder Int64-int64Decoder mDict enc nPresent bs = case enc of- PLAIN _ -> VU.convert (readNInt64 nPresent bs)- RLE_DICTIONARY _ -> unboxedLookupDict mDict nPresent bs getInt64- PLAIN_DICTIONARY _ -> unboxedLookupDict mDict nPresent bs getInt64- _ -> error ("int64Decoder: unsupported encoding " ++ show enc)- where- getInt64 (DInt64 ds) i = ds VB.! i- getInt64 d _ = error ("int64Decoder: wrong dict type, got " ++ show d)--int96Decoder :: PageDecoder UTCTime-int96Decoder mDict enc nPresent bs = case enc of- PLAIN _ -> VB.fromList (readNInt96 nPresent bs)- RLE_DICTIONARY _ -> lookupDict mDict nPresent bs getInt96- PLAIN_DICTIONARY _ -> lookupDict mDict nPresent bs getInt96- _ -> error ("int96Decoder: unsupported encoding " ++ show enc)- where- getInt96 (DInt96 ds) i = ds VB.! i- getInt96 d _ = error ("int96Decoder: wrong dict type, got " ++ show d)--floatDecoder :: UnboxedPageDecoder Float-floatDecoder mDict enc nPresent bs = case enc of- PLAIN _ -> VU.convert (readNFloat nPresent bs)- RLE_DICTIONARY _ -> unboxedLookupDict mDict nPresent bs getFloat- PLAIN_DICTIONARY _ -> unboxedLookupDict mDict nPresent bs getFloat- _ -> error ("floatDecoder: unsupported encoding " ++ show enc)- where- getFloat (DFloat ds) i = ds VB.! i- getFloat d _ = error ("floatDecoder: wrong dict type, got " ++ show d)--doubleDecoder :: UnboxedPageDecoder Double-doubleDecoder mDict enc nPresent bs = case enc of- PLAIN _ -> VU.convert (readNDouble nPresent bs)- RLE_DICTIONARY _ -> unboxedLookupDict mDict nPresent bs getDouble- PLAIN_DICTIONARY _ -> unboxedLookupDict mDict nPresent bs getDouble- _ -> error ("doubleDecoder: unsupported encoding " ++ show enc)- where- getDouble (DDouble ds) i = ds VB.! i- getDouble d _ = error ("doubleDecoder: wrong dict type, got " ++ show d)--byteArrayDecoder :: PageDecoder T.Text-byteArrayDecoder mDict enc nPresent bs = case enc of- PLAIN _ -> VB.fromList (readNTexts nPresent bs)- RLE_DICTIONARY _ -> lookupDict mDict nPresent bs getText- PLAIN_DICTIONARY _ -> lookupDict mDict nPresent bs getText- _ -> error ("byteArrayDecoder: unsupported encoding " ++ show enc)- where- getText (DText ds) i = ds VB.! i- getText d _ = error ("byteArrayDecoder: wrong dict type, got " ++ show d)--fixedLenByteArrayDecoder :: Int -> PageDecoder T.Text-fixedLenByteArrayDecoder len mDict enc nPresent bs = case enc of- PLAIN _ -> VB.fromList (readNFixedTexts len nPresent bs)- RLE_DICTIONARY _ -> lookupDict mDict nPresent bs getText- PLAIN_DICTIONARY _ -> lookupDict mDict nPresent bs getText- _ -> error ("fixedLenByteArrayDecoder: unsupported encoding " ++ show enc)- where- getText (DText ds) i = ds VB.! i- getText d _ = error ("fixedLenByteArrayDecoder: wrong dict type, got " ++ show d)--{- | Shared dictionary-path helper: decode @nPresent@ RLE/bit-packed indices-and look each one up in the dictionary.--}-lookupDict ::- Maybe DictVals ->- Int ->- BS.ByteString ->- (DictVals -> Int -> a) ->- VB.Vector a-lookupDict mDict nPresent bs f = case mDict of- Nothing -> error "Dictionary-encoded page but no dictionary page seen"- Just dict ->- let (idxs, _) = decodeDictIndices nPresent bs- in VB.generate nPresent (f dict . VU.unsafeIndex idxs)--unboxedLookupDict ::- (VU.Unbox a) =>- Maybe DictVals ->- Int ->- BS.ByteString ->- (DictVals -> Int -> a) ->- VU.Vector a-unboxedLookupDict mDict nPresent bs f = case mDict of- Nothing -> error "Dictionary-encoded page but no dictionary page seen"- Just dict ->- let (idxs, _) = decodeDictIndices nPresent bs- in VU.generate nPresent (f dict . VU.unsafeIndex idxs)---- ------------------------------------------------------------------------------ Core page-iteration loop--- ------------------------------------------------------------------------------- | Read the raw (compressed) byte range for a column chunk.-readChunkBytes ::- (RandomAccess m) =>- ColumnChunk ->- m (CompressionCodec, ThriftType, BS.ByteString)-readChunkBytes columnChunk = do- let meta = fromJust . unField $ columnChunk.cc_meta_data- codec = unField meta.cmd_codec- pType = unField meta.cmd_type- dataOffset = fromIntegral . unField $ meta.cmd_data_page_offset- dictOffset = fromIntegral <$> unField meta.cmd_dictionary_page_offset- offset = fromMaybe dataOffset dictOffset- compLen = fromIntegral . unField $ meta.cmd_total_compressed_size- rawBytes <- readBytes (Range offset compLen)- return (codec, pType, rawBytes)--{- | An 'Unfold' from a 'ColumnChunk' to per-page value triples.--The seed is a 'ColumnChunk'. The inject step reads the chunk's compressed-bytes and discovers the codec and physical type from the column metadata.-Codec and type are then threaded through the unfold state along with the-running dictionary and remaining bytes, so no intermediate list or-concatenation step is needed. Use with 'Stream.unfoldEach' to produce a-flat stream of per-page results directly from a stream of column chunks.--Dictionary pages are consumed silently and update the running dictionary-that is threaded through the unfold state.--The internal state is-@(Maybe DictVals, BS.ByteString, CompressionCodec, ThriftType)@.---- TODO: when a page index is available, use it here to compute which page--- byte ranges to request from the RandomAccess layer instead of reading the--- entire column chunk in one contiguous read.---- TODO: accept an optional row-range and use the column/offset page index--- (when present in file metadata) to Skip pages whose row range does not--- overlap the requested range, avoiding decompression of irrelevant pages--- entirely.--}-readPages ::- (RandomAccess m, MonadIO m, VG.Vector v a) =>- ColumnDescription ->- (Maybe DictVals -> Encoding -> Int -> BS.ByteString -> v a) ->- Unfold m ColumnChunk (v a, VU.Vector Int, VU.Vector Int)-readPages description decoder = mkUnfoldM step inject- where- maxDef = fromIntegral description.maxDefinitionLevel :: Int- maxRep = fromIntegral description.maxRepetitionLevel :: Int-- -- Inject: read chunk bytes; put codec and pType into state.- inject cc = do- (codec, pType, rawBytes) <- readChunkBytes cc- return (Nothing, rawBytes, codec, pType)-- step (dict, bs, codec, pType)- | BS.null bs = return Stop- | otherwise = case parsePageHeader bs of- Left e -> error ("readPages: failed to parse page header: " ++ e)- Right (rest, hdr) -> do- let compSz = fromIntegral . unField $ hdr.ph_compressed_page_size- uncmpSz = fromIntegral . unField $ hdr.ph_uncompressed_page_size- (pageData, rest') = BS.splitAt compSz rest- case unField hdr.ph_type of- DICTIONARY_PAGE _ -> do- let dictHdr =- fromMaybe- (error "DICTIONARY_PAGE: missing dictionary page header")- (unField hdr.ph_dictionary_page_header)- numVals = unField dictHdr.diph_num_values- decompressed <- liftIO $ decompressData uncmpSz codec pageData- let d = readDictVals pType decompressed numVals description.typeLength- return $ Skip (Just d, rest', codec, pType)- DATA_PAGE _ -> do- let dph =- fromMaybe- (error "DATA_PAGE: missing data page header")- (unField hdr.ph_data_page_header)- n = fromIntegral . unField $ dph.dph_num_values- enc = unField dph.dph_encoding- decompressed <- liftIO $ decompressData uncmpSz codec pageData- let (defLvls, repLvls, nPresent, valBytes) =- readLevelsV1 n maxDef maxRep decompressed- triple = (decoder dict enc nPresent valBytes, defLvls, repLvls)- return $ Yield triple (dict, rest', codec, pType)- DATA_PAGE_V2 _ -> do- let dph2 =- fromMaybe- (error "DATA_PAGE_V2: missing data page header v2")- (unField hdr.ph_data_page_header_v2)- n = fromIntegral . unField $ dph2.dph2_num_values- enc = unField dph2.dph2_encoding- defLen = unField dph2.dph2_definition_levels_byte_length- repLen = unField dph2.dph2_repetition_levels_byte_length- -- V2: levels are never compressed; only the value- -- payload is (optionally) compressed.- isCompressed = fromMaybe True (unField dph2.dph2_is_compressed)- (defLvls, repLvls, nPresent, compValBytes) =- readLevelsV2 n maxDef maxRep repLen defLen pageData- valBytes <-- if isCompressed- then liftIO $ decompressData uncmpSz codec compValBytes- else pure compValBytes- let triple = (decoder dict enc nPresent valBytes, defLvls, repLvls)- return $ Yield triple (dict, rest', codec, pType)- INDEX_PAGE _ -> return $ Skip (dict, rest', codec, pType)---- ------------------------------------------------------------------------------ Page header parsing--- -----------------------------------------------------------------------------parsePageHeader :: BS.ByteString -> Either String (BS.ByteString, PageHeader)-parsePageHeader = decodeWithLeftovers Pinch.compactProtocol---- ------------------------------------------------------------------------------ Batch value readers--- -----------------------------------------------------------------------------readNBool :: Int -> BS.ByteString -> [Bool]-readNBool count bs =- let totalBytes = (count + 7) `div` 8- bits =- concatMap- (\b -> map (\i -> (b `shiftR` i) .&. 1 == 1) [0 .. 7])- (BS.unpack (BS.take totalBytes bs))- in take count bits--readNInt32 :: Int -> BS.ByteString -> VU.Vector Int32-readNInt32 n bs = VU.generate n $ \i -> littleEndianInt32 (BS.drop (4 * i) bs)--readNInt64 :: Int -> BS.ByteString -> VU.Vector Int64-readNInt64 n bs = VU.generate n $ \i ->- fromIntegral (littleEndianWord64 (BS.drop (8 * i) bs))--readNInt96 :: Int -> BS.ByteString -> [UTCTime]-readNInt96 0 _ = []-readNInt96 n bs = int96ToUTCTime (BS.take 12 bs) : readNInt96 (n - 1) (BS.drop 12 bs)--readNFloat :: Int -> BS.ByteString -> VU.Vector Float-readNFloat n bs = VU.generate n $ \i ->- castWord32ToFloat (littleEndianWord32 (BS.drop (4 * i) bs))--readNDouble :: Int -> BS.ByteString -> VU.Vector Double-readNDouble n bs = VU.generate n $ \i ->- castWord64ToDouble (littleEndianWord64 (BS.drop (8 * i) bs))--readNTexts :: Int -> BS.ByteString -> [T.Text]-readNTexts 0 _ = []-readNTexts n bs =- let len = fromIntegral . littleEndianInt32 . BS.take 4 $ bs- text = decodeUtf8Lenient . BS.take len . BS.drop 4 $ bs- in text : readNTexts (n - 1) (BS.drop (4 + len) bs)--readNFixedTexts :: Int -> Int -> BS.ByteString -> [T.Text]-readNFixedTexts _ 0 _ = []-readNFixedTexts len n bs =- decodeUtf8Lenient (BS.take len bs)- : readNFixedTexts len (n - 1) (BS.drop len bs)
− src/DataFrame/IO/Parquet/Schema.hs
@@ -1,86 +0,0 @@-{-# 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
@@ -1,159 +0,0 @@-{- | This module contains low-level utilities around file seeking--potentially also contains all Streamly related low-level utilities.--later this module can be renamed / moved to an internal module.--}-module DataFrame.IO.Parquet.Seeking (- SeekableHandle (getSeekableHandle),- SeekMode (..),- FileBufferedOrSeekable (..),- ForceNonSeekable,- advanceBytes,- mkFileBufferedOrSeekable,- mkSeekableHandle,- readLastBytes,- seekAndReadBytes,- seekAndStreamBytes,- withFileBufferedOrSeekable,- fSeek,- fGet,-) where--import Control.Monad-import Control.Monad.IO.Class-import qualified Data.ByteString as BS-import Data.ByteString.Unsafe (unsafeDrop, unsafeTake)-import Data.IORef-import Data.Int-import Data.Word-import Streamly.Data.Stream (Stream)-import qualified Streamly.Data.Stream as S-import qualified Streamly.External.ByteString as SBS-import qualified Streamly.FileSystem.Handle as SHandle-import System.IO--{- | This handle carries a proof that it must be seekable.-Note: Handle and SeekableHandle are not thread safe, should not be-shared across threads, beaware when running parallel/concurrent code.--Not seekable:- - stdin / stdout- - pipes / FIFOs--But regular files are always seekable. Parquet fundamentally wants random-access, a non-seekable source will not support effecient access without-buffering the entire file.--}-newtype SeekableHandle = SeekableHandle {getSeekableHandle :: Handle}--{- | If we truely want to support non-seekable files, we need to also consider the case-to buffer the entire file in memory.--Not thread safe, contains mutable reference (as Handle already is).--If we need concurrent / parallel parsing or something, we need to read into ByteString-first, not sharing the same handle.--}-data FileBufferedOrSeekable- = FileBuffered !(IORef Int64) !BS.ByteString- | FileSeekable !SeekableHandle---- | Smart constructor for SeekableHandle-mkSeekableHandle :: Handle -> IO (Maybe SeekableHandle)-mkSeekableHandle h = do- seekable <- hIsSeekable h- pure $ if seekable then Just (SeekableHandle h) else Nothing---- | For testing only-type ForceNonSeekable = Maybe Bool--{- | Smart constructor for FileBufferedOrSeekable, tries to keep in the seekable case-if possible.--}-mkFileBufferedOrSeekable ::- ForceNonSeekable -> Handle -> IO FileBufferedOrSeekable-mkFileBufferedOrSeekable forceNonSeek h = do- seekable <- hIsSeekable h- if not seekable || forceNonSeek == Just True- then FileBuffered <$> newIORef 0 <*> BS.hGetContents h- else pure $ FileSeekable $ SeekableHandle h--{- | With / bracket pattern for FileBufferedOrSeekable--Warning: do not return the FileBufferedOrSeekable outside the scope of the action as-it will be closed.--}-withFileBufferedOrSeekable ::- ForceNonSeekable ->- FilePath ->- IOMode ->- (FileBufferedOrSeekable -> IO a) ->- IO a-withFileBufferedOrSeekable forceNonSeek path ioMode action = withFile path ioMode $ \h -> do- fbos <- mkFileBufferedOrSeekable forceNonSeek h- action fbos---- | Read from the end, useful for reading metadata without loading entire file-readLastBytes :: Integer -> FileBufferedOrSeekable -> IO BS.ByteString-readLastBytes n (FileSeekable sh) = do- let h = getSeekableHandle sh- hSeek h SeekFromEnd (negate n)- S.fold SBS.write (SHandle.read h)-readLastBytes n (FileBuffered i bs) = do- writeIORef i (fromIntegral $ BS.length bs)- when (n > fromIntegral (BS.length bs)) $ error "lastBytes: n > length bs"- pure $ BS.drop (BS.length bs - fromIntegral n) bs---- | Note: this does not guarantee n bytes (if it ends early)-advanceBytes :: Int -> FileBufferedOrSeekable -> IO BS.ByteString-advanceBytes = seekAndReadBytes Nothing---- | Note: this does not guarantee n bytes (if it ends early)-seekAndReadBytes ::- Maybe (SeekMode, Integer) -> Int -> FileBufferedOrSeekable -> IO BS.ByteString-seekAndReadBytes mSeek len f = seekAndStreamBytes mSeek len f >>= S.fold SBS.write--{- | Warning: the stream produced from this function accesses to the mutable handler.-if multiple streams are pulled from the same handler at the same time, chaos happen.-Make sure there is only one stream running at one time for each SeekableHandle,-and streams are not read again when they are not used anymore.--}-seekAndStreamBytes ::- (MonadIO m) =>- Maybe (SeekMode, Integer) -> Int -> FileBufferedOrSeekable -> m (Stream m Word8)-seekAndStreamBytes mSeek len f = do- liftIO $- case mSeek of- Nothing -> pure ()- Just (seekMode, seekTo) -> fSeek f seekMode seekTo- pure $ S.take len $ fRead f--fSeek :: FileBufferedOrSeekable -> SeekMode -> Integer -> IO ()-fSeek (FileSeekable (SeekableHandle h)) seekMode seekTo = hSeek h seekMode seekTo-fSeek (FileBuffered i _bs) AbsoluteSeek seekTo = writeIORef i (fromIntegral seekTo)-fSeek (FileBuffered i _bs) RelativeSeek seekTo = modifyIORef' i (+ fromIntegral seekTo)-fSeek (FileBuffered i bs) SeekFromEnd seekTo = writeIORef i (fromIntegral $ BS.length bs + fromIntegral seekTo)--fGet :: FileBufferedOrSeekable -> Int -> IO BS.ByteString-fGet (FileSeekable (SeekableHandle h)) n = BS.hGet h n-fGet (FileBuffered iRef bs) n- | n == 0 = pure BS.empty- | n > 0 = do- i <- fromIntegral <$> readIORef iRef- if (BS.length bs - i) < n- then if i <= BS.length bs then pure $ unsafeDrop i bs else pure BS.empty- else pure . unsafeTake n . unsafeDrop i $ bs- | otherwise = error "Can't read a negative number of bytes"--fRead :: (MonadIO m) => FileBufferedOrSeekable -> Stream m Word8-fRead (FileSeekable (SeekableHandle h)) = SHandle.read h-fRead (FileBuffered i bs) = S.concatEffect $ do- pos <- liftIO $ readIORef i- pure $- S.mapM- ( \x -> do- liftIO (modifyIORef' i (+ 1))- pure x- )- (S.unfold SBS.reader (BS.drop (fromIntegral pos) bs))
− src/DataFrame/IO/Parquet/Thrift.hs
@@ -1,584 +0,0 @@-{-# LANGUAGE DataKinds #-}-{-# LANGUAGE DeriveGeneric #-}-{-# LANGUAGE TypeFamilies #-}--module DataFrame.IO.Parquet.Thrift where--import Data.ByteString (ByteString)-import Data.Int (Int16, Int32, Int64, Int8)-import Data.Text (Text)-import GHC.Generics (Generic)-import GHC.TypeLits (KnownNat)-import Pinch (Enumeration, Field, Pinchable (..))-import qualified Pinch---- Primitive Parquet Types--- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L32-data ThriftType- = BOOLEAN (Enumeration 0)- | INT32 (Enumeration 1)- | INT64 (Enumeration 2)- | INT96 (Enumeration 3)- | FLOAT (Enumeration 4)- | DOUBLE (Enumeration 5)- | BYTE_ARRAY (Enumeration 6)- | FIXED_LEN_BYTE_ARRAY (Enumeration 7)- deriving (Eq, Show, Generic)--instance Pinchable ThriftType---- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L183-data FieldRepetitionType- = REQUIRED (Enumeration 0)- | OPTIONAL (Enumeration 1)- | REPEATED (Enumeration 2)- deriving (Eq, Show, Generic)--instance Pinchable FieldRepetitionType---- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L203-data Encoding- = PLAIN (Enumeration 0)- | -- GROUP_VAR_INT Encoding was never used- -- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L578- PLAIN_DICTIONARY (Enumeration 2)- | RLE (Enumeration 3)- | BIT_PACKED (Enumeration 4)- | DELTA_BINARY_PACKED (Enumeration 5)- | DELTA_LENGTH_BYTE_ARRAY (Enumeration 6)- | DELTA_BYTE_ARRAY (Enumeration 7)- | RLE_DICTIONARY (Enumeration 8)- | BYTE_STREAM_SPLIT (Enumeration 9)- deriving (Eq, Show, Generic)--instance Pinchable Encoding---- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L244-data CompressionCodec- = UNCOMPRESSED (Enumeration 0)- | SNAPPY (Enumeration 1)- | GZIP (Enumeration 2)- | LZO (Enumeration 3)- | BROTLI (Enumeration 4)- | LZ4 (Enumeration 5)- | ZSTD (Enumeration 6)- | LZ4_RAW (Enumeration 7)- deriving (Eq, Show, Generic)--instance Pinchable CompressionCodec---- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L261-data PageType- = DATA_PAGE (Enumeration 0)- | INDEX_PAGE (Enumeration 1)- | DICTIONARY_PAGE (Enumeration 2)- | DATA_PAGE_V2 (Enumeration 3)- deriving (Eq, Show, Generic)--instance Pinchable PageType---- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L271-data BoundaryOrder- = UNORDERED (Enumeration 0)- | ASCENDING (Enumeration 1)- | DESCENDING (Enumeration 2)- deriving (Eq, Show, Generic)--instance Pinchable BoundaryOrder---- Logical type annotations--- Empty structs can't use deriving Generic with Pinch, so we use a unit-like workaround.--- We represent empty structs as a newtype over () with a manual Pinchable instance.---- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L283--- struct StringType {}-data StringType = StringType deriving (Eq, Show)-instance Pinchable StringType where- type Tag StringType = Pinch.TStruct- pinch _ = Pinch.struct []- unpinch _ = pure StringType--data UUIDType = UUIDType deriving (Eq, Show)-instance Pinchable UUIDType where- type Tag UUIDType = Pinch.TStruct- pinch _ = Pinch.struct []- unpinch _ = pure UUIDType--data MapType = MapType deriving (Eq, Show)-instance Pinchable MapType where- type Tag MapType = Pinch.TStruct- pinch _ = Pinch.struct []- unpinch _ = pure MapType--data ListType = ListType deriving (Eq, Show)-instance Pinchable ListType where- type Tag ListType = Pinch.TStruct- pinch _ = Pinch.struct []- unpinch _ = pure ListType--data EnumType = EnumType deriving (Eq, Show)-instance Pinchable EnumType where- type Tag EnumType = Pinch.TStruct- pinch _ = Pinch.struct []- unpinch _ = pure EnumType--data DateType = DateType deriving (Eq, Show)-instance Pinchable DateType where- type Tag DateType = Pinch.TStruct- pinch _ = Pinch.struct []- unpinch _ = pure DateType--data Float16Type = Float16Type deriving (Eq, Show)-instance Pinchable Float16Type where- type Tag Float16Type = Pinch.TStruct- pinch _ = Pinch.struct []- unpinch _ = pure Float16Type--data NullType = NullType deriving (Eq, Show)-instance Pinchable NullType where- type Tag NullType = Pinch.TStruct- pinch _ = Pinch.struct []- unpinch _ = pure NullType--data JsonType = JsonType deriving (Eq, Show)-instance Pinchable JsonType where- type Tag JsonType = Pinch.TStruct- pinch _ = Pinch.struct []- unpinch _ = pure JsonType--data BsonType = BsonType deriving (Eq, Show)-instance Pinchable BsonType where- type Tag BsonType = Pinch.TStruct- pinch _ = Pinch.struct []- unpinch _ = pure BsonType--data VariantType = VariantType deriving (Eq, Show)-instance Pinchable VariantType where- type Tag VariantType = Pinch.TStruct- pinch _ = Pinch.struct []- unpinch _ = pure VariantType---- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L290-data TimeUnit- = MILLIS (Field 1 MilliSeconds)- | MICROS (Field 2 MicroSeconds)- | NANOS (Field 3 NanoSeconds)- deriving (Eq, Show, Generic)--instance Pinchable TimeUnit--data MilliSeconds = MilliSeconds deriving (Eq, Show)-instance Pinchable MilliSeconds where- type Tag MilliSeconds = Pinch.TStruct- pinch _ = Pinch.struct []- unpinch _ = pure MilliSeconds--data MicroSeconds = MicroSeconds deriving (Eq, Show)-instance Pinchable MicroSeconds where- type Tag MicroSeconds = Pinch.TStruct- pinch _ = Pinch.struct []- unpinch _ = pure MicroSeconds--data NanoSeconds = NanoSeconds deriving (Eq, Show)-instance Pinchable NanoSeconds where- type Tag NanoSeconds = Pinch.TStruct- pinch _ = Pinch.struct []- unpinch _ = pure NanoSeconds---- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L317-data DecimalType- = DecimalType- { decimal_scale :: Field 1 Int32- , decimal_precision :: Field 2 Int32- }- deriving (Eq, Show, Generic)--instance Pinchable DecimalType---- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L328-data IntType- = IntType- { int_bitWidth :: Field 1 Int8- , int_isSigned :: Field 2 Bool- }- deriving (Eq, Show, Generic)--instance Pinchable IntType---- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L338-data TimeType- = TimeType- { time_isAdjustedToUTC :: Field 1 Bool- , time_unit :: Field 2 TimeUnit- }- deriving (Eq, Show, Generic)--instance Pinchable TimeType---- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L349-data TimestampType- = TimestampType- { timestamp_isAdjustedToUTC :: Field 1 Bool- , timestamp_unit :: Field 2 TimeUnit- }- deriving (Eq, Show, Generic)--instance Pinchable TimestampType---- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L360--- union LogicalType-data LogicalType- = LT_STRING (Field 1 StringType)- | LT_MAP (Field 2 MapType)- | LT_LIST (Field 3 ListType)- | LT_ENUM (Field 4 EnumType)- | LT_DECIMAL (Field 5 DecimalType)- | LT_DATE (Field 6 DateType)- | LT_TIME (Field 7 TimeType)- | LT_TIMESTAMP (Field 8 TimestampType)- | LT_INTEGER (Field 10 IntType)- | LT_NULL (Field 11 NullType)- | LT_JSON (Field 12 JsonType)- | LT_BSON (Field 13 BsonType)- | LT_UUID (Field 14 UUIDType)- | LT_FLOAT16 (Field 15 Float16Type)- | LT_VARIANT (Field 16 VariantType)- deriving (Eq, Show, Generic)--instance Pinchable LogicalType---- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L270-data ConvertedType- = UTF8 (Enumeration 0)- | MAP (Enumeration 1)- | MAP_KEY_VALUE (Enumeration 2)- | LIST (Enumeration 3)- | ENUM (Enumeration 4)- | DECIMAL (Enumeration 5)- | DATE (Enumeration 6)- | TIME_MILLIS (Enumeration 7)- | TIME_MICROS (Enumeration 8)- | TIMESTAMP_MILLIS (Enumeration 9)- | TIMESTAMP_MICROS (Enumeration 10)- | UINT_8 (Enumeration 11)- | UINT_16 (Enumeration 12)- | UINT_32 (Enumeration 13)- | UINT_64 (Enumeration 14)- | INT_8 (Enumeration 15)- | INT_16 (Enumeration 16)- | INT_32 (Enumeration 17)- | INT_64 (Enumeration 18)- | JSON (Enumeration 19)- | BSON (Enumeration 20)- | INTERVAL (Enumeration 21)- deriving (Eq, Show, Generic)--instance Pinchable ConvertedType---- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L505-data SchemaElement- = SchemaElement- { schematype :: Field 1 (Maybe ThriftType) -- called just type in parquet.thrift- , type_length :: Field 2 (Maybe Int32)- , repetition_type :: Field 3 (Maybe FieldRepetitionType)- , name :: Field 4 Text- , num_children :: Field 5 (Maybe Int32)- , converted_type :: Field 6 (Maybe ConvertedType)- , scale :: Field 7 (Maybe Int32)- , precision :: Field 8 (Maybe Int32)- , field_id :: Field 9 (Maybe Int32)- , logicalType :: Field 10 (Maybe LogicalType)- }- deriving (Eq, Show, Generic)--instance Pinchable SchemaElement---- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L560-data Statistics- = Statistics- { stats_max :: Field 1 (Maybe ByteString)- , stats_min :: Field 2 (Maybe ByteString)- , stats_null_count :: Field 3 (Maybe Int64)- , stats_distinct_count :: Field 4 (Maybe Int64)- , stats_max_value :: Field 5 (Maybe ByteString)- , stats_min_value :: Field 6 (Maybe ByteString)- , stats_is_max_value_exact :: Field 7 (Maybe Bool)- , stats_is_min_value_exact :: Field 8 (Maybe Bool)- }- deriving (Eq, Show, Generic)--instance Pinchable Statistics---- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L600-data PageEncodingStats- = PageEncodingStats- { pes_page_type :: Field 1 PageType- , pes_encoding :: Field 2 Encoding- , pes_count :: Field 3 Int32- }- deriving (Eq, Show, Generic)--instance Pinchable PageEncodingStats---- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L614-data ColumnMetaData- = ColumnMetaData- { cmd_type :: Field 1 ThriftType- , cmd_encodings :: Field 2 [Encoding]- , cmd_path_in_schema :: Field 3 [Text]- , cmd_codec :: Field 4 CompressionCodec- , cmd_num_values :: Field 5 Int64- , cmd_total_uncompressed_size :: Field 6 Int64- , cmd_total_compressed_size :: Field 7 Int64- , cmd_key_value_metadata :: Field 8 (Maybe [KeyValue])- , cmd_data_page_offset :: Field 9 Int64- , cmd_index_page_offset :: Field 10 (Maybe Int64)- , cmd_dictionary_page_offset :: Field 11 (Maybe Int64)- , cmd_statistics :: Field 12 (Maybe Statistics)- , cmd_encoding_stats :: Field 13 (Maybe [PageEncodingStats])- , cmd_bloom_filter_offset :: Field 14 (Maybe Int64)- , cmd_bloom_filter_length :: Field 15 (Maybe Int32)- }- deriving (Eq, Show, Generic)--instance Pinchable ColumnMetaData---- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L875-data EncryptionWithFooterKey = EncryptionWithFooterKey deriving (Eq, Show)-instance Pinchable EncryptionWithFooterKey where- type Tag EncryptionWithFooterKey = Pinch.TStruct- pinch _ = Pinch.struct []- unpinch _ = pure EncryptionWithFooterKey---- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L883-data EncryptionWithColumnKey- = EncryptionWithColumnKey- { ewck_path_in_schema :: Field 1 [Text]- , ewck_key_metadata :: Field 2 (Maybe ByteString)- }- deriving (Eq, Show, Generic)--instance Pinchable EncryptionWithColumnKey---- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L893--- union ColumnCryptoMetaData-data ColumnCryptoMetaData- = CCM_ENCRYPTION_WITH_FOOTER_KEY (Field 1 EncryptionWithFooterKey)- | CCM_ENCRYPTION_WITH_COLUMN_KEY (Field 2 EncryptionWithColumnKey)- deriving (Eq, Show, Generic)--instance Pinchable ColumnCryptoMetaData---- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L899-data ColumnChunk- = ColumnChunk- { cc_file_path :: Field 1 (Maybe Text)- , cc_file_offset :: Field 2 Int64- , cc_meta_data :: Field 3 (Maybe ColumnMetaData)- , cc_offset_index_offset :: Field 4 (Maybe Int64)- , cc_offset_index_length :: Field 5 (Maybe Int32)- , cc_column_index_offset :: Field 6 (Maybe Int64)- , cc_column_index_length :: Field 7 (Maybe Int32)- , cc_crypto_metadata :: Field 8 (Maybe ColumnCryptoMetaData)- , cc_encrypted_column_metadata :: Field 9 (Maybe ByteString)- }- deriving (Eq, Show, Generic)--instance Pinchable ColumnChunk---- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L940-data SortingColumn- = SortingColumn- { sc_column_idx :: Field 1 Int32- , sc_descending :: Field 2 Bool- , sc_nulls_first :: Field 3 Bool- }- deriving (Eq, Show, Generic)--instance Pinchable SortingColumn---- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L958-data RowGroup- = RowGroup- { rg_columns :: Field 1 [ColumnChunk]- , rg_total_byte_size :: Field 2 Int64- , rg_num_rows :: Field 3 Int64- , rg_sorting_columns :: Field 4 (Maybe [SortingColumn])- , rg_file_offset :: Field 5 (Maybe Int64)- , rg_total_compressed_size :: Field 6 (Maybe Int64)- , rg_ordinal :: Field 7 (Maybe Int16)- }- deriving (Eq, Show, Generic)--instance Pinchable RowGroup---- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L980-data KeyValue- = KeyValue- { kv_key :: Field 1 Text- , kv_value :: Field 2 (Maybe Text)- }- deriving (Eq, Show, Generic)--instance Pinchable KeyValue---- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L990--- union ColumnOrder-newtype ColumnOrder- = TYPE_ORDER (Field 1 TypeDefinedOrder)- deriving (Eq, Show, Generic)--instance Pinchable ColumnOrder---- Empty struct for TYPE_ORDER-data TypeDefinedOrder = TypeDefinedOrder deriving (Eq, Show)-instance Pinchable TypeDefinedOrder where- type Tag TypeDefinedOrder = Pinch.TStruct- pinch _ = Pinch.struct []- unpinch _ = pure TypeDefinedOrder---- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L1094-data AesGcmV1- = AesGcmV1- { aes_gcm_v1_aad_prefix :: Field 1 (Maybe ByteString)- , aes_gcm_v1_aad_file_unique :: Field 2 (Maybe ByteString)- , aes_gcm_v1_supply_aad_prefix :: Field 3 (Maybe Bool)- }- deriving (Eq, Show, Generic)--instance Pinchable AesGcmV1---- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L1107-data AesGcmCtrV1- = AesGcmCtrV1- { aes_gcm_ctr_v1_aad_prefix :: Field 1 (Maybe ByteString)- , aes_gcm_ctr_v1_aad_file_unique :: Field 2 (Maybe ByteString)- , aes_gcm_ctr_v1_supply_aad_prefix :: Field 3 (Maybe Bool)- }- deriving (Eq, Show, Generic)--instance Pinchable AesGcmCtrV1---- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L1118--- union EncryptionAlgorithm-data EncryptionAlgorithm- = AES_GCM_V1 (Field 1 AesGcmV1)- | AES_GCM_CTR_V1 (Field 2 AesGcmCtrV1)- deriving (Eq, Show, Generic)--instance Pinchable EncryptionAlgorithm---- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L1001-data PageLocation- = PageLocation- { pl_offset :: Field 1 Int64- , pl_compressed_page_size :: Field 2 Int32- , pl_first_row_index :: Field 3 Int64- }- deriving (Eq, Show, Generic)--instance Pinchable PageLocation---- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L1017-data OffsetIndex- = OffsetIndex- { oi_page_locations :: Field 1 [PageLocation]- , oi_unencoded_byte_array_data_bytes :: Field 2 (Maybe [Int64])- }- deriving (Eq, Show, Generic)--instance Pinchable OffsetIndex---- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L1033-data ColumnIndex- = ColumnIndex- { ci_null_pages :: Field 1 [Bool]- , ci_min_values :: Field 2 [ByteString]- , ci_max_values :: Field 3 [ByteString]- , ci_boundary_order :: Field 4 BoundaryOrder- , ci_null_counts :: Field 5 (Maybe [Int64])- , ci_repetition_level_histograms :: Field 6 (Maybe [Int64])- , ci_definition_level_histograms :: Field 7 (Maybe [Int64])- }- deriving (Eq, Show, Generic)--instance Pinchable ColumnIndex---- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L1248-data DataPageHeader- = DataPageHeader- { dph_num_values :: Field 1 Int32- , dph_encoding :: Field 2 Encoding- , dph_definition_level_encoding :: Field 3 Encoding- , dph_repetition_level_encoding :: Field 4 Encoding- , dph_statistics :: Field 5 (Maybe Statistics)- }- deriving (Eq, Show, Generic)--instance Pinchable DataPageHeader--data IndexPageHeader = IndexPageHeader deriving (Eq, Show)-instance Pinchable IndexPageHeader where- type Tag IndexPageHeader = Pinch.TStruct- pinch _ = Pinch.struct []- unpinch _ = pure IndexPageHeader--data DictionaryPageHeader- = DictionaryPageHeader- { diph_num_values :: Field 1 Int32- , diph_encoding :: Field 2 Encoding- , diph_is_sorted :: Field 3 (Maybe Bool)- }- deriving (Eq, Show, Generic)--instance Pinchable DictionaryPageHeader--data DataPageHeaderV2- = DataPageHeaderV2- { dph2_num_values :: Field 1 Int32- , dph2_num_nulls :: Field 2 Int32- , dph2_num_rows :: Field 3 Int32- , dph2_encoding :: Field 4 Encoding- , dph2_definition_levels_byte_length :: Field 5 Int32- , dph2_repetition_levels_byte_length :: Field 6 Int32- , dph2_is_compressed :: Field 7 (Maybe Bool)- , dph2_statistics :: Field 8 (Maybe Statistics)- }- deriving (Eq, Show, Generic)--instance Pinchable DataPageHeaderV2--data PageHeader- = PageHeader- { ph_type :: Field 1 PageType- , ph_uncompressed_page_size :: Field 2 Int32- , ph_compressed_page_size :: Field 3 Int32- , ph_crc :: Field 4 (Maybe Int32)- , ph_data_page_header :: Field 5 (Maybe DataPageHeader)- , ph_index_page_header :: Field 6 (Maybe IndexPageHeader)- , ph_dictionary_page_header :: Field 7 (Maybe DictionaryPageHeader)- , ph_data_page_header_v2 :: Field 8 (Maybe DataPageHeaderV2)- }- deriving (Eq, Show, Generic)--instance Pinchable PageHeader---- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L1277-data FileMetadata- = FileMetadata- { version :: Field 1 Int32- , schema :: Field 2 [SchemaElement]- , num_rows :: Field 3 Int64- , row_groups :: Field 4 [RowGroup]- , key_value_metadata :: Field 5 (Maybe [KeyValue])- , created_by :: Field 6 (Maybe Text)- , column_orders :: Field 7 (Maybe [ColumnOrder])- , encryption_algorithm :: Field 8 (Maybe EncryptionAlgorithm)- , footer_signing_key_metadata :: Field 9 (Maybe ByteString)- }- deriving (Eq, Show, Generic)--instance Pinchable FileMetadata--unField :: (KnownNat n) => Field n a -> a-unField (Pinch.Field a) = a
− src/DataFrame/IO/Parquet/Time.hs
@@ -1,67 +0,0 @@-{-# LANGUAGE NumericUnderscores #-}--module DataFrame.IO.Parquet.Time where--import qualified Data.ByteString as BS-import Data.Time-import Data.Word--import DataFrame.Internal.Binary (- littleEndianWord32,- littleEndianWord64,- word32ToLittleEndian,- word64ToLittleEndian,- )--int96ToUTCTime :: BS.ByteString -> UTCTime-int96ToUTCTime bytes- | BS.length bytes /= 12 = error "INT96 must be exactly 12 bytes"- | otherwise =- let (nanosBytes, julianBytes) = BS.splitAt 8 bytes- nanosSinceMidnight = littleEndianWord64 nanosBytes- julianDay = littleEndianWord32 julianBytes- in julianDayAndNanosToUTCTime (fromIntegral julianDay) nanosSinceMidnight--julianDayAndNanosToUTCTime :: Integer -> Word64 -> UTCTime-julianDayAndNanosToUTCTime julianDay nanosSinceMidnight =- let day = julianDayToDay julianDay- secondsSinceMidnight = fromIntegral nanosSinceMidnight / (1_000_000_000 :: Double)- diffTime = secondsToDiffTime (floor secondsSinceMidnight)- in UTCTime day diffTime--julianDayToDay :: Integer -> Day-julianDayToDay julianDay =- let a = julianDay + 32_044- b = (4 * a + 3) `div` 146_097- c = a - (146_097 * b) `div` 4- d = (4 * c + 3) `div` 1461- e = c - (1461 * d) `div` 4- m = (5 * e + 2) `div` 153- day = e - (153 * m + 2) `div` 5 + 1- month = m + 3 - 12 * (m `div` 10)- year = 100 * b + d - 4800 + m `div` 10- in fromGregorian year (fromIntegral month) (fromIntegral day)---- I include this here even though it's unused because we'll likely use--- it for the writer. Since int96 is deprecated this is only included for completeness anyway.-utcTimeToInt96 :: UTCTime -> BS.ByteString-utcTimeToInt96 (UTCTime day diffTime) =- let julianDay = dayToJulianDay day- nanosSinceMidnight = floor (realToFrac diffTime * (1_000_000_000 :: Double))- nanosBytes = word64ToLittleEndian nanosSinceMidnight- julianBytes = word32ToLittleEndian (fromIntegral julianDay)- in nanosBytes `BS.append` julianBytes--dayToJulianDay :: Day -> Integer-dayToJulianDay day =- let (year, month, dayOfMonth) = toGregorian day- a = (fromIntegral $ (14 - fromIntegral month) `div` (12 :: Integer)) :: Integer- y = fromIntegral $ year + 4800 - a- m = fromIntegral $ month + 12 * fromIntegral a - 3- in fromIntegral dayOfMonth- + (153 * m + 2) `div` 5- + 365 * y- + y `div` 4- - y `div` 100- + y `div` 400- - 32_045
− src/DataFrame/IO/Parquet/Utils.hs
@@ -1,384 +0,0 @@-{-# 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
@@ -1,78 +0,0 @@-{-# 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/Binary.hs
@@ -1,94 +0,0 @@-{-# LANGUAGE BangPatterns #-}--module DataFrame.Internal.Binary where--import Data.Bits (Bits (unsafeShiftL, (.|.)))-import Data.ByteString (toStrict)-import qualified Data.ByteString as BS-import Data.ByteString.Builder (toLazyByteString, word32LE, word64LE)-import qualified Data.ByteString.Unsafe as BS-import Data.Int (Int32)-import Data.Word (Word32, Word64, Word8)--littleEndianWord32 :: BS.ByteString -> Word32-littleEndianWord32 bytes- | len >= 4 =- assembleWord32- (BS.unsafeIndex bytes 0)- (BS.unsafeIndex bytes 1)- (BS.unsafeIndex bytes 2)- (BS.unsafeIndex bytes 3)- | otherwise =- assembleWord32- (byteAtOrZero len bytes 0)- (byteAtOrZero len bytes 1)- (byteAtOrZero len bytes 2)- (byteAtOrZero len bytes 3)- where- len = BS.length bytes-{-# INLINE littleEndianWord32 #-}--littleEndianWord64 :: BS.ByteString -> Word64-littleEndianWord64 bytes- | len >= 8 =- assembleWord64- (BS.index bytes 0)- (BS.index bytes 1)- (BS.index bytes 2)- (BS.index bytes 3)- (BS.index bytes 4)- (BS.index bytes 5)- (BS.index bytes 6)- (BS.index bytes 7)- | otherwise =- assembleWord64- (byteAtOrZero len bytes 0)- (byteAtOrZero len bytes 1)- (byteAtOrZero len bytes 2)- (byteAtOrZero len bytes 3)- (byteAtOrZero len bytes 4)- (byteAtOrZero len bytes 5)- (byteAtOrZero len bytes 6)- (byteAtOrZero len bytes 7)- where- len = BS.length bytes-{-# INLINE littleEndianWord64 #-}--littleEndianInt32 :: BS.ByteString -> Int32-littleEndianInt32 = fromIntegral . littleEndianWord32-{-# INLINE littleEndianInt32 #-}--word64ToLittleEndian :: Word64 -> BS.ByteString-word64ToLittleEndian = toStrict . toLazyByteString . word64LE-{-# INLINE word64ToLittleEndian #-}--word32ToLittleEndian :: Word32 -> BS.ByteString-word32ToLittleEndian = toStrict . toLazyByteString . word32LE-{-# INLINE word32ToLittleEndian #-}--byteAtOrZero :: Int -> BS.ByteString -> Int -> Word8-byteAtOrZero len bytes i- | i >= 0 && i < len = BS.unsafeIndex bytes i- | otherwise = 0-{-# INLINE byteAtOrZero #-}--assembleWord32 :: Word8 -> Word8 -> Word8 -> Word8 -> Word32-assembleWord32 !b0 !b1 !b2 !b3 =- fromIntegral b0- .|. (fromIntegral b1 `unsafeShiftL` 8)- .|. (fromIntegral b2 `unsafeShiftL` 16)- .|. (fromIntegral b3 `unsafeShiftL` 24)-{-# INLINE assembleWord32 #-}--assembleWord64 ::- Word8 -> Word8 -> Word8 -> Word8 -> Word8 -> Word8 -> Word8 -> Word8 -> Word64-assembleWord64 !b0 !b1 !b2 !b3 !b4 !b5 !b6 !b7 =- fromIntegral b0- .|. (fromIntegral b1 `unsafeShiftL` 8)- .|. (fromIntegral b2 `unsafeShiftL` 16)- .|. (fromIntegral b3 `unsafeShiftL` 24)- .|. (fromIntegral b4 `unsafeShiftL` 32)- .|. (fromIntegral b5 `unsafeShiftL` 40)- .|. (fromIntegral b6 `unsafeShiftL` 48)- .|. (fromIntegral b7 `unsafeShiftL` 56)-{-# INLINE assembleWord64 #-}
− src/DataFrame/Internal/Column.hs
@@ -1,1761 +0,0 @@-{-# LANGUAGE AllowAmbiguousTypes #-}-{-# LANGUAGE BangPatterns #-}-{-# LANGUAGE ConstraintKinds #-}-{-# LANGUAGE DataKinds #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE FlexibleInstances #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE InstanceSigs #-}-{-# LANGUAGE LambdaCase #-}-{-# LANGUAGE MultiParamTypeClasses #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE PolyKinds #-}-{-# LANGUAGE RankNTypes #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}-{-# LANGUAGE TypeFamilies #-}-{-# LANGUAGE UndecidableInstances #-}--module DataFrame.Internal.Column where--import qualified Data.Text as T-import qualified Data.Vector as VB-import qualified Data.Vector.Generic as VG-import qualified Data.Vector.Mutable as VBM-import qualified Data.Vector.Unboxed as VU-import qualified Data.Vector.Unboxed.Mutable as VUM--import Control.DeepSeq (NFData (..), rnf)-import Control.Exception (throw)-import Control.Monad (forM_, when)-import Control.Monad.ST (ST, runST)-import Data.Bits (- complement,- popCount,- setBit,- shiftL,- shiftR,- testBit,- (.&.),- )-import Data.Kind (Type)-import Data.Maybe-import Data.These-import Data.Type.Equality (TestEquality (..))-import Data.Word (Word8)-import DataFrame.Errors-import DataFrame.Internal.Parsing-import DataFrame.Internal.Types-import System.IO.Unsafe (unsafePerformIO)-import System.Random-import Type.Reflection---- | A bit-packed validity bitmap. Bit @i@ = 1 means row @i@ is valid (not null).-type Bitmap = VU.Vector Word8--{- | Our representation of a column is a GADT that can store data based on the underlying data.--This allows us to pattern match on data kinds and limit some operations to only some-kinds of vectors. Nullability is represented via an optional bit-packed 'Bitmap':-@Nothing@ = no nulls; @Just bm@ = bit @i@ of @bm@ is 1 iff row @i@ is valid.--}-data Column where- BoxedColumn :: (Columnable a) => Maybe Bitmap -> VB.Vector a -> Column- UnboxedColumn ::- (Columnable a, VU.Unbox a) => Maybe Bitmap -> VU.Vector a -> Column--{- | A mutable companion struct to dataframe columns.--Used mostly as an intermediate structure for I/O.--}-data MutableColumn where- MBoxedColumn :: (Columnable a) => VBM.IOVector a -> MutableColumn- MUnboxedColumn :: (Columnable a, VU.Unbox a) => VUM.IOVector a -> MutableColumn---- ------------------------------------------------------------------------------ Bitmap helpers--- ------------------------------------------------------------------------------- | Test whether row @i@ is valid (not null) in a bitmap.-bitmapTestBit :: Bitmap -> Int -> Bool-bitmapTestBit bm i = testBit (VU.unsafeIndex bm (i `shiftR` 3)) (i .&. 7)-{-# INLINE bitmapTestBit #-}---- | Build a fully-valid bitmap for @n@ rows (all bits set).-allValidBitmap :: Int -> Bitmap-allValidBitmap n =- let bytes = (n + 7) `shiftR` 3- lastBits = n .&. 7- full = VU.replicate (bytes - 1) 0xFF- lastByte = if lastBits == 0 then 0xFF else (1 `shiftL` lastBits) - 1- in if bytes == 0 then VU.empty else VU.snoc full lastByte-{-# INLINE allValidBitmap #-}--{- | Build a bitmap from a @VU.Vector Word8@ validity vector-(1 = valid, 0 = null), as produced by Arrow / Parquet decoders.--}-buildBitmapFromValid :: VU.Vector Word8 -> Bitmap-buildBitmapFromValid valid =- let n = VU.length valid- bytes = (n + 7) `shiftR` 3- in VU.generate bytes $ \b ->- let base = b `shiftL` 3- setBitIf acc bit =- let idx = base + bit- in if idx < n && VU.unsafeIndex valid idx /= 0- then setBit acc bit- else acc- in foldl setBitIf (0 :: Word8) [0 .. 7]--{- | Build a bitmap from a list of null-row indices.-@nullIdxs@ are the positions that are NULL.--}-buildBitmapFromNulls :: Int -> [Int] -> Bitmap-buildBitmapFromNulls n nullIdxs =- let base = allValidBitmap n- in VU.modify- ( \mv ->- forM_ nullIdxs $ \i -> do- let byteIdx = i `shiftR` 3- bitIdx = i .&. 7- v <- VUM.unsafeRead mv byteIdx- VUM.unsafeWrite mv byteIdx (clearBit8 v bitIdx)- )- base- where- clearBit8 :: Word8 -> Int -> Word8- clearBit8 b bit = b .&. complement (1 `shiftL` bit)---- | Slice a bitmap for rows @[start .. start+len-1]@.-bitmapSlice :: Int -> Int -> Bitmap -> Bitmap-bitmapSlice start len bm- | start .&. 7 == 0 =- -- byte-aligned: simple slice; clamp so we never ask for more bytes than exist- let startByte = start `shiftR` 3- bytes = min ((len + 7) `shiftR` 3) (VU.length bm - startByte)- in VU.slice startByte bytes bm- | otherwise =- -- non-aligned: unpack bit-by-bit and repack- let n = min len (VU.length bm `shiftL` 3 - start)- in buildBitmapFromValid $- VU.generate n $- \i -> if bitmapTestBit bm (start + i) then 1 else 0---- | Concatenate two bitmaps covering @n1@ and @n2@ rows respectively.-bitmapConcat :: Int -> Bitmap -> Int -> Bitmap -> Bitmap-bitmapConcat n1 bm1 n2 bm2 =- buildBitmapFromValid $- VU.generate (n1 + n2) $ \i ->- if i < n1- then if bitmapTestBit bm1 i then 1 else 0- else if bitmapTestBit bm2 (i - n1) then 1 else 0---- | Combine two bitmaps with AND (both must be valid for result to be valid).-mergeBitmaps :: Bitmap -> Bitmap -> Bitmap-mergeBitmaps = VU.zipWith (.&.)--{- | Materialize a nullable column from @VB.Vector (Maybe a)@.-When @a@ is unboxable, creates an 'UnboxedColumn' (more compact).-Otherwise creates a 'BoxedColumn'.-Always attaches a bitmap so the column is recognized as nullable even when-no 'Nothing' values are present (preserves the Maybe type marker).--}-fromMaybeVec :: forall a. (Columnable a) => VB.Vector (Maybe a) -> Column-fromMaybeVec v = case sUnbox @a of- STrue -> fromMaybeVecUnboxed v- SFalse ->- let n = VB.length v- nullIdxs = [i | i <- [0 .. n - 1], isNothing (VB.unsafeIndex v i)]- bm = if null nullIdxs then allValidBitmap n else buildBitmapFromNulls n nullIdxs- dat = VB.map (fromMaybe (errorWithoutStackTrace "fromMaybeVec: Nothing slot")) v- in BoxedColumn (Just bm) dat--{- | Materialize a nullable 'UnboxedColumn' to @VB.Vector (Maybe a)@ using runST.-Always attaches a bitmap so the column is recognized as nullable even when-no 'Nothing' values are present (preserves the Maybe type marker).--}-fromMaybeVecUnboxed ::- forall a. (Columnable a, VU.Unbox a) => VB.Vector (Maybe a) -> Column-fromMaybeVecUnboxed v =- let n = VB.length v- nullIdxs = [i | i <- [0 .. n - 1], isNothing (VB.unsafeIndex v i)]- bm = if null nullIdxs then allValidBitmap n else buildBitmapFromNulls n nullIdxs- dat = runST $ do- mv <- VUM.new n- VG.iforM_ v $ \i mx -> forM_ mx (VUM.unsafeWrite mv i)- VU.unsafeFreeze mv- in UnboxedColumn (Just bm) dat---- | Materialize an element from a column at index @i@, respecting the bitmap.-columnElemIsNull :: Column -> Int -> Bool-columnElemIsNull (BoxedColumn (Just bm) _) i = not (bitmapTestBit bm i)-columnElemIsNull (UnboxedColumn (Just bm) _) i = not (bitmapTestBit bm i)-columnElemIsNull _ _ = False---- | Return the 'Maybe Bitmap' from a column.-columnBitmap :: Column -> Maybe Bitmap-columnBitmap (BoxedColumn bm _) = bm-columnBitmap (UnboxedColumn bm _) = bm---- ------------------------------------------------------------------------------ End bitmap helpers--- -----------------------------------------------------------------------------{- | A TypedColumn is a wrapper around our type-erased column.-It is used to type check expressions on columns.--Note: there is no guarantee that the Phanton type is the-same as the underlying vector type.--}-data TypedColumn a where- TColumn :: (Columnable a) => Column -> TypedColumn a--instance (Eq a) => Eq (TypedColumn a) where- (==) :: (Eq a) => TypedColumn a -> TypedColumn a -> Bool- (==) (TColumn a) (TColumn b) = a == b---- | Gets the underlying value from a TypedColumn.-unwrapTypedColumn :: TypedColumn a -> Column-unwrapTypedColumn (TColumn value) = value---- | Gets the underlying vector from a TypedColumn.-vectorFromTypedColumn :: TypedColumn a -> VB.Vector a-vectorFromTypedColumn (TColumn value) = either throw id (toVector value)---- | Checks if a column contains missing values (has a bitmap).-hasMissing :: Column -> Bool-hasMissing (BoxedColumn (Just _) _) = True-hasMissing (UnboxedColumn (Just _) _) = True-hasMissing _ = False---- | Checks if a column contains only missing values.-allMissing :: Column -> Bool-allMissing (BoxedColumn (Just bm) col) = VU.all (== 0) bm && not (VB.null col)-allMissing (UnboxedColumn (Just bm) col) = VU.all (== 0) bm && not (VU.null col)-allMissing _ = False---- | Checks if a column contains numeric values.-isNumeric :: Column -> Bool-isNumeric (UnboxedColumn _ (_vec :: VU.Vector a)) = case sNumeric @a of- STrue -> True- _ -> False-isNumeric (BoxedColumn _ (_vec :: VB.Vector a)) = case testEquality (typeRep @a) (typeRep @Integer) of- Nothing -> False- Just Refl -> True--{- | Checks if a column is of a given type values.-For nullable columns (@BoxedColumn (Just _)@ or @UnboxedColumn (Just _)@),-also returns @True@ when @a = Maybe b@ and the column stores @b@ internally.--}-hasElemType :: forall a. (Columnable a) => Column -> Bool-hasElemType = \case- BoxedColumn bm (_column :: VB.Vector b) -> checkBoxed bm (typeRep @b)- UnboxedColumn bm (_column :: VU.Vector b) -> checkUnboxed bm (typeRep @b)- where- -- Direct type match- directMatch :: forall (b :: Type). TypeRep b -> Bool- directMatch = isJust . testEquality (typeRep @a)- -- For a nullable column (has bitmap), also accept a = Maybe b- checkMaybe :: forall (b :: Type). TypeRep b -> Bool- checkMaybe tb = case typeRep @a of- App tMaybe tInner -> case eqTypeRep tMaybe (typeRep @Maybe) of- Just HRefl -> isJust (testEquality tInner tb)- Nothing -> False- _ -> False- checkBoxed :: forall (b :: Type). Maybe Bitmap -> TypeRep b -> Bool- checkBoxed bm tb = directMatch tb || (isJust bm && checkMaybe tb)- checkUnboxed :: forall (b :: Type). Maybe Bitmap -> TypeRep b -> Bool- checkUnboxed bm tb = directMatch tb || (isJust bm && checkMaybe tb)---- | An internal/debugging function to get the column type of a column.-columnVersionString :: Column -> String-columnVersionString column = case column of- BoxedColumn Nothing _ -> "Boxed"- BoxedColumn (Just _) _ -> "NullableBoxed"- UnboxedColumn Nothing _ -> "Unboxed"- UnboxedColumn (Just _) _ -> "NullableUnboxed"--{- | An internal/debugging function to get the type stored in the outermost vector-of a column.--}-columnTypeString :: Column -> String-columnTypeString column = case column of- BoxedColumn Nothing (_ :: VB.Vector a) -> show (typeRep @a)- BoxedColumn (Just _) (_ :: VB.Vector a) -> showMaybeType @a- UnboxedColumn Nothing (_ :: VU.Vector a) -> show (typeRep @a)- UnboxedColumn (Just _) (_ :: VU.Vector a) -> showMaybeType @a- where- showMaybeType :: forall a. (Typeable a) => String- showMaybeType =- let s = show (typeRep @a)- in "Maybe " ++ if ' ' `elem` s then "(" ++ s ++ ")" else s--instance (Show a) => Show (TypedColumn a) where- show :: (Show a) => TypedColumn a -> String- show (TColumn col) = show col--instance NFData Column where- rnf (BoxedColumn Nothing (v :: VB.Vector a)) = VB.foldl' (const (`seq` ())) () v- rnf (BoxedColumn (Just bm) (v :: VB.Vector a)) =- let n = VB.length v- go !i- | i >= n = ()- | bitmapTestBit bm i = VB.unsafeIndex v i `seq` go (i + 1)- | otherwise = go (i + 1)- in go 0- rnf (UnboxedColumn _ v) = v `seq` ()--instance Show Column where- show :: Column -> String- show (BoxedColumn Nothing column) = show column- show (BoxedColumn (Just bm) column) =- let n = VB.length column- elems =- [ if bitmapTestBit bm i then show (VB.unsafeIndex column i) else "null"- | i <- [0 .. n - 1]- ]- in "[" ++ foldl (\acc e -> if null acc then e else acc ++ "," ++ e) "" elems ++ "]"- show (UnboxedColumn Nothing column) = show column- show (UnboxedColumn (Just bm) column) =- let n = VU.length column- elems =- [ if bitmapTestBit bm i then show (VU.unsafeIndex column i) else "null"- | i <- [0 .. n - 1]- ]- in "[" ++ foldl (\acc e -> if null acc then e else acc ++ "," ++ e) "" elems ++ "]"--{- | Compare two nullable boxed columns element by element, skipping null slots.-Uses a manual loop to avoid stream fusion forcing null-slot error thunks.--}-eqBoxedCols ::- (Eq a) => Maybe Bitmap -> VB.Vector a -> Maybe Bitmap -> VB.Vector a -> Bool-eqBoxedCols bm1 a bm2 b- | VB.length a /= VB.length b = False- | otherwise = go 0- where- !n = VB.length a- go !i- | i >= n = True- | nullA || nullB = (nullA == nullB) && go (i + 1)- | VB.unsafeIndex a i == VB.unsafeIndex b i = go (i + 1)- | otherwise = False- where- nullA = maybe False (\bm -> not (bitmapTestBit bm i)) bm1- nullB = maybe False (\bm -> not (bitmapTestBit bm i)) bm2-{-# INLINE eqBoxedCols #-}--instance Eq Column where- (==) :: Column -> Column -> Bool- (==) (BoxedColumn bm1 (a :: VB.Vector t1)) (BoxedColumn bm2 (b :: VB.Vector t2)) =- case testEquality (typeRep @t1) (typeRep @t2) of- Nothing -> False- Just Refl -> eqBoxedCols bm1 a bm2 b- (==) (UnboxedColumn bm1 (a :: VU.Vector t1)) (UnboxedColumn bm2 (b :: VU.Vector t2)) =- case testEquality (typeRep @t1) (typeRep @t2) of- Nothing -> False- Just Refl ->- VU.length a == VU.length b- && VU.and- ( VU.imap- ( \i x ->- let nullA = maybe False (\bm -> not (bitmapTestBit bm i)) bm1- nullB = maybe False (\bm -> not (bitmapTestBit bm i)) bm2- in if nullA || nullB then nullA == nullB else x == VU.unsafeIndex b i- )- a- )- (==) _ _ = False--{- | A class for converting a vector to a column of the appropriate type.-Given each Rep we tell the `toColumnRep` function which Column type to pick.--}-class ColumnifyRep (r :: Rep) a where- toColumnRep :: VB.Vector a -> Column---- | Constraint synonym for what we can put into columns.-type Columnable a =- ( Columnable' a- , ColumnifyRep (KindOf a) a- , UnboxIf a- , IntegralIf a- , FloatingIf a- , SBoolI (Unboxable a)- , SBoolI (Numeric a)- , SBoolI (IntegralTypes a)- , SBoolI (FloatingTypes a)- )--instance- (Columnable a, VU.Unbox a) =>- ColumnifyRep 'RUnboxed a- where- toColumnRep :: (Columnable a, VUM.Unbox a) => VB.Vector a -> Column- toColumnRep v = UnboxedColumn Nothing (VU.convert v)--instance- (Columnable a) =>- ColumnifyRep 'RBoxed a- where- toColumnRep :: (Columnable a) => VB.Vector a -> Column- toColumnRep = BoxedColumn Nothing--instance- (Columnable a) =>- ColumnifyRep 'RNullableBoxed (Maybe a)- where- toColumnRep :: (Columnable a) => VB.Vector (Maybe a) -> Column- toColumnRep = fromMaybeVec--{- | O(n) Convert a vector to a column. Automatically picks the best representation of a vector to store the underlying data in.--__Examples:__--@-> import qualified Data.Vector as V-> fromVector (VB.fromList [(1 :: Int), 2, 3, 4])-[1,2,3,4]-@--}-fromVector ::- forall a.- (Columnable a, ColumnifyRep (KindOf a) a) =>- VB.Vector a -> Column-fromVector = toColumnRep @(KindOf a)--{- | O(n) Convert an unboxed vector to a column. This avoids the extra conversion if you already have the data in an unboxed vector.--__Examples:__--@-> import qualified Data.Vector.Unboxed as V-> fromUnboxedVector (VB.fromList [(1 :: Int), 2, 3, 4])-[1,2,3,4]-@--}-fromUnboxedVector ::- forall a. (Columnable a, VU.Unbox a) => VU.Vector a -> Column-fromUnboxedVector = UnboxedColumn Nothing--{- | O(n) Convert a list to a column. Automatically picks the best representation of a vector to store the underlying data in.--__Examples:__--@-> fromList [(1 :: Int), 2, 3, 4]-[1,2,3,4]-@--}-fromList ::- forall a.- (Columnable a, ColumnifyRep (KindOf a) a) =>- [a] -> Column-fromList = toColumnRep @(KindOf a) . VB.fromList--{- | O(n) Create a column of random elements within a range.--Takes a random number generator, a length, and a lower and upper bound for the random values.--__Examples:__--@-> import System.Random (mkStdGen)-> mkRandom (mkStdGen 42) 4 0 10-[4,2,6,5]-@--}-mkRandom ::- (RandomGen g, Columnable a, ColumnifyRep (KindOf a) a, UniformRange a) =>- g -> Int -> a -> a -> Column-mkRandom pureGen k lo hi = fromList $ go pureGen k- where- go _g 0 = []- go g n =- let- (!v, !g') = uniformR (lo, hi) g- in- v : go g' (n - 1)---- An internal helper for type errors-throwTypeMismatch ::- forall (a :: Type) (b :: Type).- (Typeable a, Typeable b) => Either DataFrameException Column-throwTypeMismatch =- Left $- TypeMismatchException- MkTypeErrorContext- { userType = Right (typeRep @b)- , expectedType = Right (typeRep @a)- , callingFunctionName = Nothing- , errorColumnName = Nothing- }---- | An internal function to map a function over the values of a column.-mapColumn ::- forall b c.- (Columnable b, Columnable c) =>- (b -> c) -> Column -> Either DataFrameException Column-mapColumn f = \case- BoxedColumn bm (col :: VB.Vector a) -> runBoxed bm col- UnboxedColumn bm (col :: VU.Vector a) -> runUnboxed bm col- where- runBoxed ::- forall a.- (Columnable a) =>- Maybe Bitmap -> VB.Vector a -> Either DataFrameException Column- runBoxed bm col = case testEquality (typeRep @b) (typeRep @(Maybe a)) of- -- user maps over Maybe a (nullable column as Maybe)- Just Refl ->- let !n = VB.length col- in -- Build result directly without intermediate Maybe vector to avoid- -- fusion forcing null slots via VU.convert.- Right $ case sUnbox @c of- STrue -> UnboxedColumn Nothing $- VU.generate n $ \i ->- f- ( if maybe True (`bitmapTestBit` i) bm- then Just (VB.unsafeIndex col i)- else Nothing- )- SFalse -> fromVector @c $- VB.generate n $ \i ->- f- ( if maybe True (`bitmapTestBit` i) bm- then Just (VB.unsafeIndex col i)- else Nothing- )- Nothing -> case testEquality (typeRep @a) (typeRep @b) of- Just Refl ->- -- user maps over inner type a; preserve bitmap- Right $ case sUnbox @c of- STrue -> UnboxedColumn bm (VU.generate (VB.length col) (f . VB.unsafeIndex col))- SFalse -> BoxedColumn bm (VB.map f col)- Nothing -> throwTypeMismatch @a @b-- runUnboxed ::- forall a.- (Columnable a, VU.Unbox a) =>- Maybe Bitmap -> VU.Vector a -> Either DataFrameException Column- runUnboxed bm col = case testEquality (typeRep @b) (typeRep @(Maybe a)) of- Just Refl ->- let !n = VU.length col- in Right $ case sUnbox @c of- STrue -> UnboxedColumn Nothing $- VU.generate n $ \i ->- f- ( if maybe True (`bitmapTestBit` i) bm- then Just (VU.unsafeIndex col i)- else Nothing- )- SFalse -> fromVector @c $- VB.generate n $ \i ->- f- ( if maybe True (`bitmapTestBit` i) bm- then Just (VU.unsafeIndex col i)- else Nothing- )- Nothing -> case testEquality (typeRep @a) (typeRep @b) of- Just Refl -> Right $ case sUnbox @c of- STrue -> UnboxedColumn bm (VU.map f col)- SFalse -> BoxedColumn bm (VB.generate (VU.length col) (f . VU.unsafeIndex col))- Nothing -> throwTypeMismatch @a @b-{-# INLINEABLE mapColumn #-}---- | Applies a function that returns an unboxed result to an unboxed vector, storing the result in a column.-imapColumn ::- forall b c.- (Columnable b, Columnable c) =>- (Int -> b -> c) -> Column -> Either DataFrameException Column-imapColumn f = \case- BoxedColumn bm (col :: VB.Vector a) -> runBoxed bm col- UnboxedColumn bm (col :: VU.Vector a) -> runUnboxed bm col- where- runBoxed ::- forall a.- (Columnable a) =>- Maybe Bitmap -> VB.Vector a -> Either DataFrameException Column- runBoxed bm col = case testEquality (typeRep @a) (typeRep @b) of- Just Refl -> Right $ case sUnbox @c of- STrue ->- UnboxedColumn- bm- (VU.generate (VB.length col) (\i -> f i (VB.unsafeIndex col i)))- SFalse -> BoxedColumn bm (VB.imap f col)- Nothing -> throwTypeMismatch @a @b-- runUnboxed ::- forall a.- (Columnable a, VU.Unbox a) =>- Maybe Bitmap -> VU.Vector a -> Either DataFrameException Column- runUnboxed bm col = case testEquality (typeRep @a) (typeRep @b) of- Just Refl -> Right $ case sUnbox @c of- STrue -> UnboxedColumn bm (VU.imap f col)- SFalse -> BoxedColumn bm (VB.imap f (VG.convert col))- Nothing -> throwTypeMismatch @a @b---- | O(1) Gets the number of elements in the column.-columnLength :: Column -> Int-columnLength (BoxedColumn _ xs) = VB.length xs-columnLength (UnboxedColumn _ xs) = VU.length xs-{-# INLINE columnLength #-}---- | O(n) Gets the number of non-null elements in the column.-numElements :: Column -> Int-numElements (BoxedColumn Nothing xs) = VB.length xs-numElements (BoxedColumn (Just bm) _xs) = VU.foldl' (\acc b -> acc + popCount b) 0 bm-numElements (UnboxedColumn Nothing xs) = VU.length xs-numElements (UnboxedColumn (Just bm) _xs) = VU.foldl' (\acc b -> acc + popCount b) 0 bm-{-# INLINE numElements #-}---- | O(n) Takes the first n values of a column.-takeColumn :: Int -> Column -> Column-takeColumn n (BoxedColumn bm xs) =- BoxedColumn (fmap (bitmapSlice 0 n) bm) (VG.take n xs)-takeColumn n (UnboxedColumn bm xs) =- UnboxedColumn (fmap (bitmapSlice 0 n) bm) (VG.take n xs)-{-# INLINE takeColumn #-}---- | O(n) Takes the last n values of a column.-takeLastColumn :: Int -> Column -> Column-takeLastColumn n column = sliceColumn (columnLength column - n) n column-{-# INLINE takeLastColumn #-}---- | O(n) Takes n values after a given column index.-sliceColumn :: Int -> Int -> Column -> Column-sliceColumn start n (BoxedColumn bm xs) =- BoxedColumn (fmap (bitmapSlice start n) bm) (VG.slice start n xs)-sliceColumn start n (UnboxedColumn bm xs) =- UnboxedColumn (fmap (bitmapSlice start n) bm) (VG.slice start n xs)-{-# INLINE sliceColumn #-}---- | O(n) Selects the elements at a given set of indices. Does not change the order.-atIndicesStable :: VU.Vector Int -> Column -> Column-atIndicesStable indexes (BoxedColumn bm column) =- BoxedColumn- ( fmap- ( \bm0 ->- buildBitmapFromValid $- VU.map (\i -> if bitmapTestBit bm0 i then 1 else 0) indexes- )- bm- )- ( VB.generate- (VU.length indexes)- ((column `VB.unsafeIndex`) . (indexes `VU.unsafeIndex`))- )-atIndicesStable indexes (UnboxedColumn bm column) =- UnboxedColumn- ( fmap- ( \bm0 ->- buildBitmapFromValid $- VU.map (\i -> if bitmapTestBit bm0 i then 1 else 0) indexes- )- bm- )- (VU.unsafeBackpermute column indexes)-{-# INLINE atIndicesStable #-}--{- | Like 'atIndicesStable' but treats negative indices as null.-Keeps the index vector fully unboxed (no @VB.Vector (Maybe Int)@).--}-gatherWithSentinel :: VU.Vector Int -> Column -> Column-gatherWithSentinel indices col =- let !n = VU.length indices- newBm = buildBitmapFromValid $ VU.generate n $ \i ->- if VU.unsafeIndex indices i < 0 then 0 else 1- in case col of- BoxedColumn srcBm v ->- let dat = VB.generate n $ \i ->- let !idx = VU.unsafeIndex indices i- in if idx < 0 then VB.unsafeIndex v 0 else VB.unsafeIndex v idx- bm = case srcBm of- Nothing -> Just newBm- Just sb ->- Just- ( mergeBitmaps- newBm- ( buildBitmapFromValid $ VU.generate n $ \i ->- let idx = VU.unsafeIndex indices i- in if idx >= 0 && bitmapTestBit sb idx then 1 else 0- )- )- in BoxedColumn bm dat- UnboxedColumn srcBm v ->- let dat = runST $ do- mv <- VUM.new n- VG.iforM_ indices $ \i idx ->- when (idx >= 0) $ VUM.unsafeWrite mv i (VU.unsafeIndex v idx)- VU.unsafeFreeze mv- bm = case srcBm of- Nothing -> Just newBm- Just sb ->- Just- ( mergeBitmaps- newBm- ( buildBitmapFromValid $ VU.generate n $ \i ->- let idx = VU.unsafeIndex indices i- in if idx >= 0 && bitmapTestBit sb idx then 1 else 0- )- )- in UnboxedColumn bm dat-{-# INLINE gatherWithSentinel #-}---- | Internal helper to get indices in a boxed vector.-getIndices :: VU.Vector Int -> VB.Vector a -> VB.Vector a-getIndices indices xs = VB.generate (VU.length indices) (\i -> xs VB.! (indices VU.! i))-{-# INLINE getIndices #-}---- | Internal helper to get indices in an unboxed vector.-getIndicesUnboxed :: (VU.Unbox a) => VU.Vector Int -> VU.Vector a -> VU.Vector a-getIndicesUnboxed indices xs = VU.generate (VU.length indices) (\i -> xs VU.! (indices VU.! i))-{-# INLINE getIndicesUnboxed #-}--findIndices ::- forall a.- (Columnable a) =>- (a -> Bool) ->- Column ->- Either DataFrameException (VU.Vector Int)-findIndices predicate = \case- BoxedColumn _ (v :: VB.Vector b) -> run v VG.convert- UnboxedColumn _ (v :: VU.Vector b) -> run v id- where- run ::- forall b v.- (Typeable b, VG.Vector v b, VG.Vector v Int) =>- v b ->- (v Int -> VU.Vector Int) ->- Either DataFrameException (VU.Vector Int)- run column finalize = case testEquality (typeRep @a) (typeRep @b) of- Just Refl -> Right . finalize $ VG.findIndices predicate column- Nothing ->- Left $- TypeMismatchException- MkTypeErrorContext- { userType = Right (typeRep @a)- , expectedType = Right (typeRep @b)- , callingFunctionName = Just "findIndices"- , errorColumnName = Nothing- }---- | Fold (right) column with index.-ifoldrColumn ::- forall a b.- (Columnable a, Columnable b) =>- (Int -> a -> b -> b) -> b -> Column -> Either DataFrameException b-ifoldrColumn f acc = \case- BoxedColumn _ column -> foldrWorker column- UnboxedColumn _ column -> foldrWorker column- where- foldrWorker ::- forall c v.- (Typeable c, VG.Vector v c) =>- v c ->- Either DataFrameException b- foldrWorker vec = case testEquality (typeRep @a) (typeRep @c) of- Just Refl -> pure $ VG.ifoldr f acc vec- Nothing ->- Left $- TypeMismatchException- ( MkTypeErrorContext- { userType = Right (typeRep @a)- , expectedType = Right (typeRep @c)- , callingFunctionName = Just "ifoldrColumn"- , errorColumnName = Nothing- }- )--foldlColumn ::- forall a b.- (Columnable a, Columnable b) =>- (b -> a -> b) -> b -> Column -> Either DataFrameException b-foldlColumn f acc = \case- BoxedColumn _ column -> foldlWorker column- UnboxedColumn _ column -> foldlWorker column- where- foldlWorker ::- forall c v.- (Typeable c, VG.Vector v c) =>- v c ->- Either DataFrameException b- foldlWorker vec = case testEquality (typeRep @a) (typeRep @c) of- Just Refl -> pure $ VG.foldl' f acc vec- Nothing ->- Left $- TypeMismatchException- ( MkTypeErrorContext- { userType = Right (typeRep @a)- , expectedType = Right (typeRep @c)- , callingFunctionName = Just "ifoldrColumn"- , errorColumnName = Nothing- }- )--foldl1Column ::- forall a.- (Columnable a) =>- (a -> a -> a) -> Column -> Either DataFrameException a-foldl1Column f = \case- BoxedColumn _ column -> foldl1Worker column- UnboxedColumn _ column -> foldl1Worker column- where- foldl1Worker ::- forall c v.- (Typeable c, VG.Vector v c) =>- v c ->- Either DataFrameException a- foldl1Worker vec = case testEquality (typeRep @a) (typeRep @c) of- Just Refl -> pure $ VG.foldl1' f vec- Nothing ->- Left $- TypeMismatchException- ( MkTypeErrorContext- { userType = Right (typeRep @a)- , expectedType = Right (typeRep @c)- , callingFunctionName = Just "foldl1Column"- , errorColumnName = Nothing- }- )--{- | O(n) Seedless fold over groups using the first element of each group as seed.-Like 'foldDirectGroups' but for the case where no initial accumulator is available.--}-foldl1DirectGroups ::- forall a.- (Columnable a) =>- (a -> a -> a) ->- Column ->- VU.Vector Int ->- VU.Vector Int ->- Either DataFrameException Column-foldl1DirectGroups f col valueIndices offsets- | VU.length offsets <= 1 = pure $ fromVector @a VB.empty- | otherwise = case col of- UnboxedColumn _ (vec :: VU.Vector d) -> UnboxedColumn Nothing <$> foldl1Worker vec- BoxedColumn _ (vec :: VB.Vector d) -> BoxedColumn Nothing <$> foldl1Worker vec- where- foldl1Worker ::- forall c v.- (Typeable c, VG.Vector v c) =>- v c ->- Either DataFrameException (v c)- foldl1Worker vec = case testEquality (typeRep @a) (typeRep @c) of- Just Refl ->- Right $- VG.generate (VU.length offsets - 1) foldGroup- where- foldGroup k =- let !s = VU.unsafeIndex offsets k- !e = VU.unsafeIndex offsets (k + 1)- !seed = VG.unsafeIndex vec (VU.unsafeIndex valueIndices s)- in go (s + 1) e seed- go !i !e !acc- | i >= e = acc- | otherwise =- go (i + 1) e $!- f acc (VG.unsafeIndex vec (VU.unsafeIndex valueIndices i))- Nothing ->- Left $- TypeMismatchException- MkTypeErrorContext- { userType = Right (typeRep @a)- , expectedType = Right (typeRep @c)- , callingFunctionName = Just "foldl1DirectGroups"- , errorColumnName = Nothing- }-{-# INLINEABLE foldl1DirectGroups #-}--{- | O(n) fold over groups by scanning the column LINEARLY.-rowToGroup[i] = group index for row i.-Avoids random column reads; random writes go to the accumulator array which is-small (nGroups entries) and typically cache-resident.-When @acc@ is unboxable, uses an unboxed mutable vector for the accumulator-array, eliminating pointer indirection on every read/write.--}-foldLinearGroups ::- forall b acc.- (Columnable b, Columnable acc) =>- (acc -> b -> acc) ->- acc ->- Column ->- VU.Vector Int -> -- rowToGroup (length n)- Int -> -- nGroups- Either DataFrameException Column-foldLinearGroups f seed col rowToGroup nGroups- | nGroups == 0 = Right (fromVector @acc VB.empty)- | otherwise = case col of- UnboxedColumn _ (vec :: VU.Vector d) -> foldLinearWorker vec- BoxedColumn _ (vec :: VB.Vector d) -> foldLinearWorker vec- where- foldLinearWorker ::- forall c v.- (Typeable c, VG.Vector v c) =>- v c ->- Either DataFrameException Column- foldLinearWorker vec = case testEquality (typeRep @b) (typeRep @c) of- Just Refl ->- Right $- unsafePerformIO $- runWith- ( \readAt writeAt ->- VG.iforM_ vec $ \row x -> do- let !k = VG.unsafeIndex rowToGroup row- cur <- readAt k- writeAt k $! f cur x- )- Nothing ->- Left $- TypeMismatchException- MkTypeErrorContext- { userType = Right (typeRep @b)- , expectedType = Right (typeRep @c)- , callingFunctionName = Just "foldLinearGroups"- , errorColumnName = Nothing- }-- -- \| Allocate accumulators, run the traversal, return a frozen Column.- -- When @acc@ is unboxable, uses an unboxed mutable vector (no pointer- -- indirection per read/write) and returns UnboxedColumn directly —- -- avoiding a round-trip through VB.Vector.- runWith :: ((Int -> IO acc) -> (Int -> acc -> IO ()) -> IO ()) -> IO Column- runWith body = case sUnbox @acc of- STrue -> do- accs <- VUM.replicate nGroups seed- body (VUM.unsafeRead accs) (VUM.unsafeWrite accs)- UnboxedColumn Nothing <$> VU.unsafeFreeze accs- SFalse -> do- accs <- VBM.replicate nGroups seed- body (VBM.unsafeRead accs) (VBM.unsafeWrite accs)- fromVector @acc <$> VB.unsafeFreeze accs- {-# INLINE runWith #-}-{-# INLINEABLE foldLinearGroups #-}--headColumn :: forall a. (Columnable a) => Column -> Either DataFrameException a-headColumn = \case- BoxedColumn _ col -> headWorker col- UnboxedColumn _ col -> headWorker col- where- headWorker ::- forall c v.- (Typeable c, VG.Vector v c) =>- v c ->- Either DataFrameException a- headWorker vec = case testEquality (typeRep @a) (typeRep @c) of- Just Refl ->- if VG.null vec- then Left (EmptyDataSetException "headColumn")- else pure (VG.head vec)- Nothing ->- Left $- TypeMismatchException- ( MkTypeErrorContext- { userType = Right (typeRep @a)- , expectedType = Right (typeRep @c)- , callingFunctionName = Just "headColumn"- , errorColumnName = Nothing- }- )---- | An internal, column version of zip.-zipColumns :: Column -> Column -> Column-zipColumns (BoxedColumn _ column) (BoxedColumn _ other) = BoxedColumn Nothing (VG.zip column other)-zipColumns (BoxedColumn _ column) (UnboxedColumn _ other) =- BoxedColumn- Nothing- ( VB.generate- (min (VG.length column) (VG.length other))- (\i -> (column VG.! i, other VG.! i))- )-zipColumns (UnboxedColumn _ column) (BoxedColumn _ other) =- BoxedColumn- Nothing- ( VB.generate- (min (VG.length column) (VG.length other))- (\i -> (column VG.! i, other VG.! i))- )-zipColumns (UnboxedColumn _ column) (UnboxedColumn _ other) = UnboxedColumn Nothing (VG.zip column other)-{-# INLINE zipColumns #-}---- | Merge two columns using `These`.-mergeColumns :: Column -> Column -> Column-mergeColumns colA colB = case (colA, colB) of- (BoxedColumn bmA c1, BoxedColumn bmB c2) -> case (bmA, bmB) of- (Just ba, Just bb) ->- BoxedColumn Nothing $ mkVec c1 c2 $ \i v1 v2 ->- let nullA = not (bitmapTestBit ba i)- nullB = not (bitmapTestBit bb i)- in case (nullA, nullB) of- (True, True) -> error "mergeColumns: both null"- (False, True) -> This v1- (True, False) -> That v2- (False, False) -> These v1 v2- (Just ba, Nothing) ->- BoxedColumn Nothing $ mkVec c1 c2 $ \i v1 v2 ->- if not (bitmapTestBit ba i) then That v2 else These v1 v2- (Nothing, Just bb) ->- BoxedColumn Nothing $ mkVec c1 c2 $ \i v1 v2 ->- if not (bitmapTestBit bb i) then This v1 else These v1 v2- (Nothing, Nothing) ->- BoxedColumn Nothing $ mkVecSimple c1 c2 These- (BoxedColumn _ c1, UnboxedColumn _ c2) ->- BoxedColumn Nothing $ mkVecSimple c1 c2 These- (UnboxedColumn _ c1, BoxedColumn _ c2) ->- BoxedColumn Nothing $ mkVecSimple c1 c2 These- (UnboxedColumn _ c1, UnboxedColumn _ c2) ->- BoxedColumn Nothing $ mkVecSimple c1 c2 These- where- mkVec c1 c2 combineElements =- VB.generate- (min (VG.length c1) (VG.length c2))- (\i -> combineElements i (c1 VG.! i) (c2 VG.! i))- {-# INLINE mkVec #-}-- mkVecSimple c1 c2 f =- VB.generate- (min (VG.length c1) (VG.length c2))- (\i -> f (c1 VG.! i) (c2 VG.! i))- {-# INLINE mkVecSimple #-}-{-# INLINE mergeColumns #-}---- | An internal, column version of zipWith.-zipWithColumns ::- forall a b c.- (Columnable a, Columnable b, Columnable c) =>- (a -> b -> c) -> Column -> Column -> Either DataFrameException Column-zipWithColumns f (UnboxedColumn bmL (column :: VU.Vector d)) (UnboxedColumn bmR (other :: VU.Vector e)) = case testEquality (typeRep @a) (typeRep @d) of- Just Refl -> case testEquality (typeRep @b) (typeRep @e) of- Just Refl- -- Fast path: both plain unboxed, no bitmaps involved in the output type- | isNothing bmL- , isNothing bmR ->- pure $ case sUnbox @c of- STrue -> UnboxedColumn Nothing (VU.zipWith f column other)- SFalse -> fromVector $ VB.zipWith f (VG.convert column) (VG.convert other)- -- Type mismatch or bitmap involvement: fall through to general toVector path- _ -> zipWithColumnsGeneral f (UnboxedColumn bmL column) (UnboxedColumn bmR other)- Nothing -> zipWithColumnsGeneral f (UnboxedColumn bmL column) (UnboxedColumn bmR other)--- TODO: mchavinda - reuse pattern from interpret where we augment the--- error at the end.-zipWithColumns f left right = zipWithColumnsGeneral f left right--zipWithColumnsGeneral ::- forall a b c.- (Columnable a, Columnable b, Columnable c) =>- (a -> b -> c) -> Column -> Column -> Either DataFrameException Column-zipWithColumnsGeneral f left right = case toVector @a left of- Left (TypeMismatchException context) ->- Left $- TypeMismatchException (context{callingFunctionName = Just "zipWithColumns"})- Left e -> Left e- Right left' -> case toVector @b right of- Left (TypeMismatchException context) ->- Left $- TypeMismatchException (context{callingFunctionName = Just "zipWithColumns"})- Left e -> Left e- Right right' -> pure $ fromVector $ VB.zipWith f left' right'-{-# INLINE zipWithColumnsGeneral #-}-{-# INLINE zipWithColumns #-}---- Functions for mutable columns (intended for IO).-writeColumn :: Int -> T.Text -> MutableColumn -> IO (Either T.Text Bool)-writeColumn i value (MBoxedColumn (col :: VBM.IOVector a)) =- case testEquality (typeRep @a) (typeRep @T.Text) of- Just Refl ->- if isNullish value- then VBM.unsafeWrite col i "" >> return (Left $! value)- else VBM.unsafeWrite col i value >> return (Right True)- Nothing -> return (Left value)-writeColumn i value (MUnboxedColumn (col :: VUM.IOVector a)) =- case testEquality (typeRep @a) (typeRep @Int) of- Just Refl -> case readInt value of- Just v -> VUM.unsafeWrite col i v >> return (Right True)- Nothing -> VUM.unsafeWrite col i 0 >> return (Left value)- Nothing -> case testEquality (typeRep @a) (typeRep @Double) of- Nothing -> return (Left $! value)- Just Refl -> case readDouble value of- Just v -> VUM.unsafeWrite col i v >> return (Right True)- Nothing -> VUM.unsafeWrite col i 0 >> return (Left $! value)-{-# INLINE writeColumn #-}--freezeColumn' :: [(Int, T.Text)] -> MutableColumn -> IO Column-freezeColumn' nulls (MBoxedColumn col)- | null nulls = BoxedColumn Nothing <$> VB.unsafeFreeze col- | all (isNullish . snd) nulls = do- frozen <- VB.unsafeFreeze col- let n = VB.length frozen- bm = buildBitmapFromNulls n (map fst nulls)- return $ BoxedColumn (Just bm) frozen- | otherwise =- BoxedColumn Nothing- . VB.imap- ( \i v ->- if i `elem` map fst nulls- then Left (fromMaybe (error "UNEXPECTED ERROR DURING FREEZE") (lookup i nulls))- else Right v- )- <$> VB.unsafeFreeze col-freezeColumn' nulls (MUnboxedColumn col)- | null nulls = UnboxedColumn Nothing <$> VU.unsafeFreeze col- | all (isNullish . snd) nulls = do- c <- VU.unsafeFreeze col- let n = VU.length c- bm = buildBitmapFromNulls n (map fst nulls)- return $ UnboxedColumn (Just bm) c- | otherwise = do- c <- VU.unsafeFreeze col- return $- BoxedColumn Nothing $- VB.generate- (VU.length c)- ( \i ->- if i `elem` map fst nulls- then Left (fromMaybe (error "UNEXPECTED ERROR DURING FREEZE") (lookup i nulls))- else Right (c VU.! i)- )-{-# INLINE freezeColumn' #-}--{- | Freeze a mutable column into an @Either Text a@ column: every recorded-null position becomes @Left rawText@ (preserving the original input), every-other position becomes @Right v@. Used by CSV readers under 'EitherRead' mode.--}-freezeColumnEither :: [(Int, T.Text)] -> MutableColumn -> IO Column-freezeColumnEither nulls (MBoxedColumn col) = do- frozen <- VB.unsafeFreeze col- let nullMap = nulls- pure $- BoxedColumn Nothing $- VB.imap- ( \i v -> case lookup i nullMap of- Just t -> Left t- Nothing -> Right v- )- frozen-freezeColumnEither nulls (MUnboxedColumn col) = do- c <- VU.unsafeFreeze col- let nullMap = nulls- pure $- BoxedColumn Nothing $- VB.generate (VU.length c) $ \i ->- case lookup i nullMap of- Just t -> Left t- Nothing -> Right (c VU.! i)-{-# INLINE freezeColumnEither #-}--{- | Promote a non-nullable column to a nullable one (add an all-valid bitmap).-No-op when already nullable.--}-ensureOptional :: Column -> Column-ensureOptional c@(BoxedColumn (Just _) _) = c-ensureOptional (BoxedColumn Nothing col) =- BoxedColumn (Just (allValidBitmap (VB.length col))) col-ensureOptional c@(UnboxedColumn (Just _) _) = c-ensureOptional (UnboxedColumn Nothing col) =- UnboxedColumn (Just (allValidBitmap (VU.length col))) col---- | Fills the end of a column, up to n, with null rows. Does nothing if column has length >= n.-expandColumn :: Int -> Column -> Column-expandColumn n column@(BoxedColumn bm col)- | n <= VG.length col = column- | otherwise =- let extra = n - VG.length col- newBm = case bm of- Nothing -> Just (buildBitmapFromNulls n [VG.length col .. n - 1])- Just b ->- Just- (bitmapConcat (VG.length col) b extra (VU.replicate ((extra + 7) `shiftR` 3) 0))- -- pad data with default (undefined slot, protected by bitmap)- newCol = col <> VB.replicate extra (errorWithoutStackTrace "expandColumn: null slot")- in BoxedColumn newBm newCol-expandColumn n column@(UnboxedColumn bm col)- | n <= VG.length col = column- | otherwise =- let extra = n - VG.length col- newBm = case bm of- Nothing -> Just (buildBitmapFromNulls n [VG.length col .. n - 1])- Just b ->- Just- (bitmapConcat (VG.length col) b extra (VU.replicate ((extra + 7) `shiftR` 3) 0))- newCol = runST $ do- mv <- VUM.new n- VU.imapM_ (VUM.unsafeWrite mv) col- VU.unsafeFreeze mv- in UnboxedColumn newBm newCol---- | Fills the beginning of a column, up to n, with null rows. Does nothing if column has length >= n.-leftExpandColumn :: Int -> Column -> Column-leftExpandColumn n column@(BoxedColumn bm col)- | n <= VG.length col = column- | otherwise =- let extra = n - VG.length col- origLen = VG.length col- newBm = case bm of- Nothing -> Just (buildBitmapFromNulls n [0 .. extra - 1])- Just b ->- let nullPart = VU.replicate ((extra + 7) `shiftR` 3) 0- in Just (bitmapConcat extra nullPart origLen b)- newCol =- VB.replicate extra (errorWithoutStackTrace "leftExpandColumn: null slot") <> col- in BoxedColumn newBm newCol-leftExpandColumn n column@(UnboxedColumn bm col)- | n <= VG.length col = column- | otherwise =- let extra = n - VG.length col- origLen = VG.length col- newBm = case bm of- Nothing -> Just (buildBitmapFromNulls n [0 .. extra - 1])- Just b ->- let nullPart = VU.replicate ((extra + 7) `shiftR` 3) 0- in Just (bitmapConcat extra nullPart origLen b)- newCol = runST $ do- mv <- VUM.new n- VU.imapM_ (\i x -> VUM.unsafeWrite mv (extra + i) x) col- VU.unsafeFreeze mv- in UnboxedColumn newBm newCol--{- | Concatenates two columns.-Returns Nothing if the columns are of different types.--}-concatColumns :: Column -> Column -> Either DataFrameException Column-concatColumns left right = case (left, right) of- (BoxedColumn bmL l, BoxedColumn bmR r) -> case testEquality (typeOf l) (typeOf r) of- Just Refl ->- let newBm = case (bmL, bmR) of- (Nothing, Nothing) -> Nothing- (Just bl, Nothing) ->- Just- (bitmapConcat (VB.length l) bl (VB.length r) (allValidBitmap (VB.length r)))- (Nothing, Just br) ->- Just- (bitmapConcat (VB.length l) (allValidBitmap (VB.length l)) (VB.length r) br)- (Just bl, Just br) -> Just (bitmapConcat (VB.length l) bl (VB.length r) br)- in pure (BoxedColumn newBm (l <> r))- Nothing -> Left (mismatchErr (typeOf r) (typeOf l))- (UnboxedColumn bmL l, UnboxedColumn bmR r) -> case testEquality (typeOf l) (typeOf r) of- Just Refl ->- let newBm = case (bmL, bmR) of- (Nothing, Nothing) -> Nothing- (Just bl, Nothing) ->- Just- (bitmapConcat (VU.length l) bl (VU.length r) (allValidBitmap (VU.length r)))- (Nothing, Just br) ->- Just- (bitmapConcat (VU.length l) (allValidBitmap (VU.length l)) (VU.length r) br)- (Just bl, Just br) -> Just (bitmapConcat (VU.length l) bl (VU.length r) br)- in pure (UnboxedColumn newBm (l <> r))- Nothing -> Left (mismatchErr (typeOf r) (typeOf l))- _ -> Left (mismatchErr (typeOf right) (typeOf left))- where- mismatchErr ::- forall (x :: Type) (y :: Type). TypeRep x -> TypeRep y -> DataFrameException- mismatchErr ta tb =- withTypeable ta $- withTypeable tb $- TypeMismatchException- ( MkTypeErrorContext- { userType = Right ta- , expectedType = Right tb- , callingFunctionName = Just "concatColumns"- , errorColumnName = Nothing- }- )--{- | Concatenates two columns.--Works similar to 'concatColumns', but unlike that function, it will also combine columns of different types-by wrapping the values in an Either.--E.g. combining Column containing [1,2] with Column containing ["a","b"]-will result in a Column containing [Left 1, Left 2, Right "a", Right "b"].--}--{- | O(n) Concatenate a list of same-type columns in a single allocation.-All columns must have the same constructor and element type (as they will-within a single Parquet column). Calls 'error' on mismatch.--}-concatManyColumns :: [Column] -> Column-concatManyColumns [] = fromList ([] :: [Maybe Int])-concatManyColumns [c] = c-concatManyColumns (c0 : cs) = case c0 of- BoxedColumn bm0 v0 ->- let getCol (BoxedColumn bm v) = case testEquality (typeOf v0) (typeOf v) of- Just Refl -> (bm, v)- Nothing -> error "concatManyColumns: BoxedColumn type mismatch"- getCol _ = error "concatManyColumns: column constructor mismatch"- rest = map getCol cs- allVecs = v0 : map snd rest- allBms = bm0 : map fst rest- newBm- | all isNothing allBms = Nothing- | otherwise =- let pairs = zip allVecs allBms- expandedBms = map (\(v, mb) -> fromMaybe (allValidBitmap (VB.length v)) mb) pairs- go b1 n1 b2 n2 = bitmapConcat n1 b1 n2 b2- concatBms [] = VU.empty- concatBms [(b, _v)] = b- concatBms ((b1, v1) : (b2, v2) : rest') =- let merged = go b1 (VB.length v1) b2 (VB.length v2)- in concatBms ((merged, v1 <> v2) : rest')- in Just $ concatBms (zip expandedBms allVecs)- in BoxedColumn newBm (VB.concat allVecs)- UnboxedColumn bm0 v0 ->- let getCol (UnboxedColumn bm v) = case testEquality (typeOf v0) (typeOf v) of- Just Refl -> (bm, v)- Nothing -> error "concatManyColumns: UnboxedColumn type mismatch"- getCol _ = error "concatManyColumns: column constructor mismatch"- rest = map getCol cs- allVecs = v0 : map snd rest- allBms = bm0 : map fst rest- newBm- | all isNothing allBms = Nothing- | otherwise =- let pairs = zip allVecs allBms- expandedBms = map (\(v, mb) -> fromMaybe (allValidBitmap (VU.length v)) mb) pairs- go b1 n1 b2 n2 = bitmapConcat n1 b1 n2 b2- concatBms [] = VU.empty- concatBms [(b, _)] = b- concatBms ((b1, v1) : (b2, v2) : rest') =- let merged = go b1 (VU.length v1) b2 (VU.length v2)- in concatBms ((merged, v1 <> v2) : rest')- in Just $ concatBms (zip expandedBms allVecs)- in UnboxedColumn newBm (VU.concat allVecs)--concatColumnsEither :: Column -> Column -> Column-concatColumnsEither (BoxedColumn bmL left) (BoxedColumn bmR right) = case testEquality (typeOf left) (typeOf right) of- Nothing ->- BoxedColumn Nothing $ fmap Left left <> fmap Right right- Just Refl ->- let newBm = case (bmL, bmR) of- (Nothing, Nothing) -> Nothing- (Just bl, Nothing) ->- Just- ( bitmapConcat- (VB.length left)- bl- (VB.length right)- (allValidBitmap (VB.length right))- )- (Nothing, Just br) ->- Just- ( bitmapConcat- (VB.length left)- (allValidBitmap (VB.length left))- (VB.length right)- br- )- (Just bl, Just br) -> Just (bitmapConcat (VB.length left) bl (VB.length right) br)- in BoxedColumn newBm $ left <> right-concatColumnsEither (UnboxedColumn bmL left) (UnboxedColumn bmR right) = case testEquality (typeOf left) (typeOf right) of- Nothing ->- BoxedColumn Nothing $- fmap Left (VG.convert left) <> fmap Right (VG.convert right)- Just Refl ->- let newBm = case (bmL, bmR) of- (Nothing, Nothing) -> Nothing- (Just bl, Nothing) ->- Just- ( bitmapConcat- (VU.length left)- bl- (VU.length right)- (allValidBitmap (VU.length right))- )- (Nothing, Just br) ->- Just- ( bitmapConcat- (VU.length left)- (allValidBitmap (VU.length left))- (VU.length right)- br- )- (Just bl, Just br) -> Just (bitmapConcat (VU.length left) bl (VU.length right) br)- in UnboxedColumn newBm $ left <> right-concatColumnsEither (BoxedColumn _ left) (UnboxedColumn _ right) =- BoxedColumn Nothing $ fmap Left left <> fmap Right (VG.convert right)-concatColumnsEither (UnboxedColumn _ left) (BoxedColumn _ right) =- BoxedColumn Nothing $ fmap Left (VG.convert left) <> fmap Right right---- | Allocate a mutable column of size @n@ matching the constructor/type of the given column.-newMutableColumn :: Int -> Column -> IO MutableColumn-newMutableColumn n (BoxedColumn _ (_ :: VB.Vector a)) =- MBoxedColumn <$> (VBM.new n :: IO (VBM.IOVector a))-newMutableColumn n (UnboxedColumn _ (_ :: VU.Vector a)) =- MUnboxedColumn <$> (VUM.new n :: IO (VUM.IOVector a))---- | Copy a column chunk into a mutable column starting at offset @off@.-copyIntoMutableColumn :: MutableColumn -> Int -> Column -> IO ()-copyIntoMutableColumn (MBoxedColumn (mv :: VBM.IOVector b)) off (BoxedColumn _ (v :: VB.Vector a)) =- case testEquality (typeRep @a) (typeRep @b) of- Just Refl -> VG.imapM_ (\i x -> VBM.unsafeWrite mv (off + i) x) v- Nothing -> error "copyIntoMutableColumn: Boxed type mismatch"-copyIntoMutableColumn (MUnboxedColumn (mv :: VUM.IOVector b)) off (UnboxedColumn _ (v :: VU.Vector a)) =- case testEquality (typeRep @a) (typeRep @b) of- Just Refl -> VG.imapM_ (\i x -> VUM.unsafeWrite mv (off + i) x) v- Nothing -> error "copyIntoMutableColumn: Unboxed type mismatch"-copyIntoMutableColumn _ _ _ =- error "copyIntoMutableColumn: constructor mismatch"---- | Freeze a mutable column into an immutable column.-freezeMutableColumn :: MutableColumn -> IO Column-freezeMutableColumn (MBoxedColumn mv) = BoxedColumn Nothing <$> VB.unsafeFreeze mv-freezeMutableColumn (MUnboxedColumn mv) = UnboxedColumn Nothing <$> VU.unsafeFreeze mv--{- | O(n) Converts a column to a list. Throws an exception if the wrong type is specified.--__Examples:__--@-> column = fromList [(1 :: Int), 2, 3, 4]-> toList @Int column-[1,2,3,4]-> toList @Double column-exception: ...-@--}-toList :: forall a. (Columnable a) => Column -> [a]-toList xs = case toVector @a xs of- Left err -> throw err- Right val -> VB.toList val--{- | Converts a column to a vector of a specific type.--This is a type-safe conversion that requires the column's element type-to exactly match the requested type. You must specify the desired type-via type applications.--==== __Type Parameters__--[@a@] The element type to convert to-[@v@] The vector type (e.g., 'VU.Vector', 'VB.Vector')--==== __Examples__-->>> toVector @Int @VU.Vector column-Right (unboxed vector of Ints)-->>> toVector @Text @VB.Vector column-Right (boxed vector of Text)--==== __Returns__--* 'Right' - The converted vector if types match-* 'Left' 'TypeMismatchException' - If the column's type doesn't match the requested type--==== __See also__--For numeric conversions with automatic type coercion, see 'toDoubleVector',-'toFloatVector', and 'toIntVector'.--}-toVector ::- forall a v.- (VG.Vector v a, Columnable a) => Column -> Either DataFrameException (v a)-toVector col = case col of- BoxedColumn bm (inner :: VB.Vector c) ->- -- Check if user wants Maybe c (nullable) or c directly- case testEquality (typeRep @a) (typeRep @c) of- Just Refl -> Right $ VG.convert inner- Nothing ->- -- Try: a = Maybe c- case testEquality (typeRep @a) (typeRep @(Maybe c)) of- Just Refl ->- -- Use VB.generate to avoid fusion forcing null slots- let !n = VB.length inner- maybeVec = case bm of- Nothing -> VB.generate n (Just . VB.unsafeIndex inner)- Just bitmap -> VB.generate n $ \i ->- if bitmapTestBit bitmap i then Just (VB.unsafeIndex inner i) else Nothing- in Right $ VG.convert maybeVec- Nothing ->- Left $- TypeMismatchException- ( MkTypeErrorContext- { userType = Right (typeRep @a)- , expectedType = Right (typeRep @c)- , callingFunctionName = Just "toVector"- , errorColumnName = Nothing- }- )- UnboxedColumn bm (inner :: VU.Vector c) ->- case testEquality (typeRep @a) (typeRep @c) of- Just Refl -> Right $ VG.convert inner- Nothing ->- case testEquality (typeRep @a) (typeRep @(Maybe c)) of- Just Refl ->- let maybeVec = case bm of- Nothing -> VB.generate (VU.length inner) (Just . VU.unsafeIndex inner)- Just bitmap -> VB.generate (VU.length inner) $ \i ->- if bitmapTestBit bitmap i then Just (VU.unsafeIndex inner i) else Nothing- in Right $ VG.convert maybeVec- Nothing ->- Left $- TypeMismatchException- ( MkTypeErrorContext- { userType = Right (typeRep @a)- , expectedType = Right (typeRep @c)- , callingFunctionName = Just "toVector"- , errorColumnName = Nothing- }- )---- Some common types we will use for numerical computing.--{- | Converts a column to an unboxed vector of 'Double' values.--This function performs intelligent type coercion for numeric types:--* If the column is already 'Double', returns it directly-* If the column contains other floating-point types, converts via 'realToFrac'-* If the column contains integral types, converts via 'fromIntegral' (beware of overflow if the type is `Integer`).--==== __Optional column handling__--For 'OptionalColumn' types, 'Nothing' values are converted to @NaN@ (Not a Number).-This allows optional numeric data to be represented in the resulting vector.--==== __Returns__--* 'Right' - The converted 'Double' vector-* 'Left' 'TypeMismatchException' - If the column is not numeric--}-toDoubleVector :: Column -> Either DataFrameException (VU.Vector Double)-toDoubleVector column =- case column of- UnboxedColumn bm (f :: VU.Vector a) -> case testEquality (typeRep @a) (typeRep @Double) of- Just Refl -> case bm of- Nothing -> Right f- Just bitmap -> Right $ VU.imap (\i x -> if bitmapTestBit bitmap i then x else read "NaN") f- Nothing -> case sFloating @a of- STrue ->- Right- ( VU.imap- ( \i x -> case bm of- Just bitmap | not (bitmapTestBit bitmap i) -> read "NaN"- _ -> realToFrac x- )- f- )- SFalse -> case sIntegral @a of- STrue ->- Right- ( VU.imap- ( \i x -> case bm of- Just bitmap | not (bitmapTestBit bitmap i) -> read "NaN"- _ -> fromIntegral x- )- f- )- SFalse ->- Left $- TypeMismatchException- ( MkTypeErrorContext- { userType = Right (typeRep @Double)- , expectedType = Right (typeRep @a)- , callingFunctionName = Just "toDoubleVector"- , errorColumnName = Nothing- }- )- BoxedColumn bm (f :: VB.Vector a) -> case testEquality (typeRep @a) (typeRep @Integer) of- Just Refl ->- Right- ( VB.convert $- VB.imap- ( \i x -> case bm of- Just bitmap | not (bitmapTestBit bitmap i) -> read "NaN"- _ -> fromIntegral x- )- f- )- Nothing ->- Left $- TypeMismatchException- ( MkTypeErrorContext- { userType = Right (typeRep @Double)- , expectedType = Left (columnTypeString column) :: Either String (TypeRep ())- , callingFunctionName = Just "toDoubleVector"- , errorColumnName = Nothing- }- )--{- | Converts a column to an unboxed vector of 'Float' values.--This function performs intelligent type coercion for numeric types:--* If the column is already 'Float', returns it directly-* If the column contains other floating-point types, converts via 'realToFrac'-* If the column contains integral types, converts via 'fromIntegral'-* If the column is boxed 'Integer', converts via 'fromIntegral' (beware of overflow for 64-bit integers and `Integer`)--==== __Optional column handling__--For 'OptionalColumn' types, 'Nothing' values are converted to @NaN@ (Not a Number).-This allows optional numeric data to be represented in the resulting vector.--==== __Returns__--* 'Right' - The converted 'Float' vector-* 'Left' 'TypeMismatchException' - If the column is not numeric--==== __Precision warning__--Converting from 'Double' to 'Float' may result in loss of precision.--}-toFloatVector :: Column -> Either DataFrameException (VU.Vector Float)-toFloatVector column =- case column of- UnboxedColumn bm (f :: VU.Vector a) -> case testEquality (typeRep @a) (typeRep @Float) of- Just Refl -> case bm of- Nothing -> Right f- Just bitmap -> Right $ VU.imap (\i x -> if bitmapTestBit bitmap i then x else read "NaN") f- Nothing -> case sFloating @a of- STrue ->- Right- ( VU.imap- ( \i x -> case bm of- Just bitmap | not (bitmapTestBit bitmap i) -> read "NaN"- _ -> realToFrac x- )- f- )- SFalse -> case sIntegral @a of- STrue ->- Right- ( VU.imap- ( \i x -> case bm of- Just bitmap | not (bitmapTestBit bitmap i) -> read "NaN"- _ -> fromIntegral x- )- f- )- SFalse ->- Left $- TypeMismatchException- ( MkTypeErrorContext- { userType = Right (typeRep @Float)- , expectedType = Right (typeRep @a)- , callingFunctionName = Just "toFloatVector"- , errorColumnName = Nothing- }- )- BoxedColumn bm (f :: VB.Vector a) -> case testEquality (typeRep @a) (typeRep @Integer) of- Just Refl ->- Right- ( VB.convert $- VB.imap- ( \i x -> case bm of- Just bitmap | not (bitmapTestBit bitmap i) -> read "NaN"- _ -> fromIntegral x- )- f- )- Nothing ->- Left $- TypeMismatchException- ( MkTypeErrorContext- { userType = Right (typeRep @Float)- , expectedType = Left (columnTypeString column) :: Either String (TypeRep ())- , callingFunctionName = Just "toFloatVector"- , errorColumnName = Nothing- }- )--{- | Converts a column to an unboxed vector of 'Int' values.--This function performs intelligent type coercion for numeric types:--* If the column is already 'Int', returns it directly-* If the column contains floating-point types, rounds via 'round' and converts-* If the column contains other integral types, converts via 'fromIntegral'-* If the column is boxed 'Integer', converts via 'fromIntegral'--==== __Returns__--* 'Right' - The converted 'Int' vector-* 'Left' 'TypeMismatchException' - If the column is not numeric--==== __Note__--Unlike 'toDoubleVector' and 'toFloatVector', this function does NOT support-'OptionalColumn'. Optional columns must be handled separately.--==== __Rounding behavior__--Floating-point values are rounded to the nearest integer using 'round'.-For example: 2.5 rounds to 2, 3.5 rounds to 4 (banker's rounding).--}-toIntVector :: Column -> Either DataFrameException (VU.Vector Int)-toIntVector column =- case column of- UnboxedColumn _ (f :: VU.Vector a) -> case testEquality (typeRep @a) (typeRep @Int) of- Just Refl -> Right f- Nothing -> case sFloating @a of- STrue -> Right (VU.map (round . (realToFrac :: a -> Double)) f)- SFalse -> case sIntegral @a of- STrue -> Right (VU.map fromIntegral f)- SFalse ->- Left $- TypeMismatchException- ( MkTypeErrorContext- { userType = Right (typeRep @Int)- , expectedType = Right (typeRep @a)- , callingFunctionName = Just "toIntVector"- , errorColumnName = Nothing- }- )- BoxedColumn _ (f :: VB.Vector a) -> case testEquality (typeRep @a) (typeRep @Integer) of- Just Refl -> Right (VB.convert $ VB.map fromIntegral f)- Nothing ->- Left $- TypeMismatchException- ( MkTypeErrorContext- { userType = Right (typeRep @Int)- , expectedType = Left (columnTypeString column) :: Either String (TypeRep ())- , callingFunctionName = Just "toIntVector"- , errorColumnName = Nothing- }- )--toUnboxedVector ::- forall a.- (Columnable a, VU.Unbox a) => Column -> Either DataFrameException (VU.Vector a)-toUnboxedVector column =- case column of- UnboxedColumn _ (f :: VU.Vector b) -> case testEquality (typeRep @a) (typeRep @b) of- Just Refl -> Right f- Nothing ->- Left $- TypeMismatchException- ( MkTypeErrorContext- { userType = Right (typeRep @Int)- , expectedType = Right (typeRep @a)- , callingFunctionName = Just "toUnboxedVector"- , errorColumnName = Nothing- }- )- _ ->- Left $- TypeMismatchException- ( MkTypeErrorContext- { userType = Right (typeRep @a)- , expectedType = Left (columnTypeString column) :: Either String (TypeRep ())- , callingFunctionName = Just "toUnboxedVector"- , errorColumnName = Nothing- }- )-{-# INLINE toUnboxedVector #-}---- Shared finaliser for the two parseUnboxedColumn* helpers. Freezes--- the mutable data vector, and only materialises the bitmap when the--- column actually had nulls.-{-# INLINE finalizeParseResult #-}-finalizeParseResult ::- (VU.Unbox a) =>- VUM.STVector s a ->- VUM.STVector s Word8 ->- Bool ->- ST s (Maybe (Maybe Bitmap, VU.Vector a))-finalizeParseResult values vmask anyNull- | anyNull = do- vs <- VU.unsafeFreeze values- vm <- VU.unsafeFreeze vmask- return (Just (Just (buildBitmapFromValid vm), vs))- | otherwise = do- vs <- VU.unsafeFreeze values- return (Just (Nothing, vs))
− src/DataFrame/Internal/DataFrame.hs
@@ -1,335 +0,0 @@-{-# LANGUAGE AllowAmbiguousTypes #-}-{-# LANGUAGE ExplicitNamespaces #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE InstanceSigs #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE PatternSynonyms #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}--module DataFrame.Internal.DataFrame where--import qualified Data.Map as M-import qualified Data.Text as T-import qualified Data.Vector as V-import qualified Data.Vector.Unboxed as VU--import Control.DeepSeq (NFData (..), rnf)-import Control.Exception (throw)-import Data.Function (on)-import Data.List (sortBy, (\\))-import Data.Maybe (fromMaybe)-import Data.Type.Equality (- TestEquality (testEquality),- type (:~:) (Refl),- type (:~~:) (HRefl),- )-import DataFrame.Display.Terminal.PrettyPrint-import DataFrame.Errors-import DataFrame.Internal.Column-import DataFrame.Internal.Expression-import Text.Printf-import Type.Reflection (Typeable, eqTypeRep, typeRep, pattern App)-import Prelude hiding (null)--data DataFrame = DataFrame- { columns :: V.Vector Column- {- ^ Our main data structure stores a dataframe as- a vector of columns. This improv- -}- , columnIndices :: M.Map T.Text Int- -- ^ Keeps the column names in the order they were inserted in.- , dataframeDimensions :: (Int, Int)- -- ^ (rows, columns)- , derivingExpressions :: M.Map T.Text UExpr- }--instance NFData DataFrame where- rnf (DataFrame cols idx dims _exprs) =- rnf cols `seq` rnf idx `seq` rnf dims--{- | A record that contains information about how and what-rows are grouped in the dataframe. This can only be used with-`aggregate`.--}-data GroupedDataFrame = Grouped- { fullDataframe :: DataFrame- , groupedColumns :: [T.Text]- , valueIndices :: VU.Vector Int- , offsets :: VU.Vector Int- , rowToGroup :: VU.Vector Int- {- ^ rowToGroup[i] = group index for row i. Length n (one per row).- Built once in 'groupBy'; reused by every aggregation.- -}- }--instance Show GroupedDataFrame where- show (Grouped df cols _indices _os _rtg) =- printf- "{ keyColumns: %s groupedColumns: %s }"- (show cols)- (show (M.keys (columnIndices df) \\ cols))--instance Eq GroupedDataFrame where- (==) (Grouped df cols _indices _os _rtg) (Grouped df' cols' _indices' _os' _rtg') = (df == df') && (cols == cols')--instance Eq DataFrame where- (==) :: DataFrame -> DataFrame -> Bool- a == b =- M.keys (columnIndices a) == M.keys (columnIndices b)- && foldr- ( \(name, index) acc -> acc && (columns a V.!? index == (columns b V.!? (columnIndices b M.! name)))- )- True- (M.toList $ columnIndices a)--instance Show DataFrame where- show :: DataFrame -> String- show d =- 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 = asText Markdown---- | For showing the dataframe as a string markdown in notebooks.-toMarkdown' :: DataFrame -> String-toMarkdown' = T.unpack . toMarkdown--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 =- let s = show (typeRep @a)- in "Maybe " <> if ' ' `elem` s then "(" <> s <> ")" else s- getType (BoxedColumn Nothing (_ :: V.Vector a)) = T.pack $ show (typeRep @a)- 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.- 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 (UnboxedColumn (Just bm) 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- 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-empty =- DataFrame- { columns = V.empty- , columnIndices = M.empty- , dataframeDimensions = (0, 0)- , derivingExpressions = M.empty- }--{- | Safely retrieves a column by name from the dataframe.--Returns 'Nothing' if the column does not exist.--==== __Examples__-->>> getColumn "age" df-Just (UnboxedColumn ...)-->>> getColumn "nonexistent" df-Nothing--}-getColumn :: T.Text -> DataFrame -> Maybe Column-getColumn name df- | null df = Nothing- | otherwise = do- i <- columnIndices df M.!? name- columns df V.!? i--{- | Retrieves a column by name from the dataframe, throwing an exception if not found.--This is an unsafe version of 'getColumn' that throws 'ColumnsNotFoundException'-if the column does not exist. Use this when you are certain the column exists.--==== __Throws__--* 'ColumnsNotFoundException' - if the column with the given name does not exist--}-unsafeGetColumn :: T.Text -> DataFrame -> Column-unsafeGetColumn name df = case getColumn name df of- Nothing -> throw $ ColumnsNotFoundException [name] "" (M.keys $ columnIndices df)- Just col -> col--{- | Checks if the dataframe is empty (has no columns).--Returns 'True' if the dataframe has no columns, 'False' otherwise.-Note that a dataframe with columns but no rows is not considered null.--}-null :: DataFrame -> Bool-null df = V.null (columns df)---- | Convert a DataFrame to a CSV (comma-separated) text.-toCsv :: DataFrame -> T.Text-toCsv = toSeparated ','---- | Convert a DataFrame to a CSV (comma-separated) string.-toCsv' :: DataFrame -> String-toCsv' = T.unpack . toSeparated ','---- | Convert a DataFrame to a text representation with a custom separator.-toSeparated :: Char -> DataFrame -> T.Text-toSeparated sep df- | null df = T.empty- | otherwise =- let (rows, _) = dataframeDimensions df- headers = map fst (sortBy (compare `on` snd) (M.toList (columnIndices df)))- sepText = T.singleton sep- headerLine = T.intercalate sepText headers- dataLines = map (T.intercalate sepText . getRowAsText df) [0 .. rows - 1]- in T.unlines (headerLine : dataLines)--getRowAsText :: DataFrame -> Int -> [T.Text]-getRowAsText df i = map (`showElement` i) (V.toList (columns df))--showElement :: Column -> Int -> T.Text-showElement (BoxedColumn _ (c :: V.Vector a)) i = case c V.!? i of- Nothing -> error $ "Column index out of bounds at row " ++ show i- Just e- | Just Refl <- testEquality (typeRep @a) (typeRep @T.Text) -> e- | App t1 t2 <- typeRep @a- , Just HRefl <- eqTypeRep t1 (typeRep @Maybe) ->- case testEquality t2 (typeRep @T.Text) of- Just Refl -> fromMaybe "null" e- Nothing -> stripJust (T.pack (show e))- | otherwise -> T.pack (show e)-showElement (UnboxedColumn _ c) i = case c VU.!? i of- Nothing -> error $ "Column index out of bounds at row " ++ show i- Just e -> T.pack (show e)--stripJust :: T.Text -> T.Text-stripJust = fromMaybe "null" . T.stripPrefix "Just "
− src/DataFrame/Internal/Expression.hs
@@ -1,397 +0,0 @@-{-# LANGUAGE AllowAmbiguousTypes #-}-{-# LANGUAGE ExplicitNamespaces #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE FlexibleInstances #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE InstanceSigs #-}-{-# LANGUAGE MultiParamTypeClasses #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE RankNTypes #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}-{-# LANGUAGE UndecidableInstances #-}--module DataFrame.Internal.Expression where--import Data.String-import qualified Data.Text as T-import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl))-import qualified Data.Vector.Generic as VG-import DataFrame.Internal.Column-import Type.Reflection (Typeable, typeOf, typeRep)--data UnaryOp a b = MkUnaryOp- { unaryFn :: a -> b- , unaryName :: T.Text- , unarySymbol :: Maybe T.Text- }--data BinaryOp a b c = MkBinaryOp- { binaryFn :: a -> b -> c- , binaryName :: T.Text- , binarySymbol :: Maybe T.Text- , binaryCommutative :: Bool- , binaryPrecedence :: Int- }--data MeanAcc = MeanAcc {-# UNPACK #-} !Double {-# UNPACK #-} !Int- deriving (Show, Eq, Ord, Read)--data AggStrategy a b where- CollectAgg ::- (VG.Vector v b, Typeable v) => T.Text -> (v b -> a) -> AggStrategy a b- FoldAgg :: T.Text -> Maybe a -> (a -> b -> a) -> AggStrategy a b- MergeAgg ::- (Columnable acc) =>- T.Text ->- acc ->- (acc -> b -> acc) ->- (acc -> acc -> acc) ->- (acc -> a) ->- AggStrategy a b--data Expr a where- Col :: (Columnable a) => T.Text -> Expr a- CastWith ::- (Columnable a, Columnable b, Read a) =>- T.Text ->- T.Text ->- (Either String a -> b) ->- Expr b- CastExprWith ::- (Columnable a, Columnable b, Columnable src, Read a) =>- T.Text ->- (Either String a -> b) ->- Expr src ->- Expr b- Lit :: (Columnable a) => a -> Expr a- Unary ::- (Columnable a, Columnable b) => UnaryOp b a -> Expr b -> Expr a- Binary ::- (Columnable c, Columnable b, Columnable a) =>- BinaryOp c b a -> Expr c -> Expr b -> Expr a- If :: (Columnable a) => Expr Bool -> Expr a -> Expr a -> Expr a- Agg :: (Columnable a, Columnable b) => AggStrategy a b -> Expr b -> Expr a- Over :: (Columnable a) => [T.Text] -> Expr a -> Expr a--data UExpr where- UExpr :: (Columnable a) => Expr a -> UExpr--instance Show UExpr where- show :: UExpr -> String- show (UExpr expr) = show expr--type NamedExpr = (T.Text, UExpr)--instance (Num a, Columnable a) => Num (Expr a) where- (+) :: Expr a -> Expr a -> Expr a- (+) =- Binary- ( MkBinaryOp- { binaryFn = (+)- , binaryName = "add"- , binarySymbol = Just "+"- , binaryCommutative = True- , binaryPrecedence = 6- }- )-- (-) :: Expr a -> Expr a -> Expr a- (-) =- Binary- ( MkBinaryOp- { binaryFn = (-)- , binaryName = "sub"- , binarySymbol = Just "-"- , binaryCommutative = False- , binaryPrecedence = 6- }- )-- (*) :: Expr a -> Expr a -> Expr a- (*) =- Binary- ( MkBinaryOp- { binaryFn = (*)- , binaryName = "mult"- , binarySymbol = Just "*"- , binaryCommutative = True- , binaryPrecedence = 7- }- )-- fromInteger :: Integer -> Expr a- fromInteger = Lit . fromInteger-- negate :: Expr a -> Expr a- negate =- Unary- (MkUnaryOp{unaryFn = negate, unaryName = "negate", unarySymbol = Nothing})-- abs :: (Num a) => Expr a -> Expr a- abs = Unary (MkUnaryOp{unaryFn = abs, unaryName = "abs", unarySymbol = Nothing})-- signum :: (Num a) => Expr a -> Expr a- signum =- Unary- (MkUnaryOp{unaryFn = signum, unaryName = "signum", unarySymbol = Nothing})--add :: (Num a, Columnable a) => Expr a -> Expr a -> Expr a-add = (+)--sub :: (Num a, Columnable a) => Expr a -> Expr a -> Expr a-sub = (-)--mult :: (Num a, Columnable a) => Expr a -> Expr a -> Expr a-mult = (*)--instance (Fractional a, Columnable a) => Fractional (Expr a) where- fromRational :: (Fractional a, Columnable a) => Rational -> Expr a- fromRational = Lit . fromRational-- (/) :: (Fractional a, Columnable a) => Expr a -> Expr a -> Expr a- (/) =- Binary- ( MkBinaryOp- { binaryFn = (/)- , binaryName = "divide"- , binarySymbol = Just "/"- , binaryCommutative = False- , binaryPrecedence = 7- }- )--divide :: (Fractional a, Columnable a) => Expr a -> Expr a -> Expr a-divide = (/)--instance (IsString a, Columnable a) => IsString (Expr a) where- fromString :: String -> Expr a- fromString s = Lit (fromString s)--instance (Floating a, Columnable a) => Floating (Expr a) where- pi :: (Floating a, Columnable a) => Expr a- pi = Lit pi- exp :: (Floating a, Columnable a) => Expr a -> Expr a- exp = Unary (MkUnaryOp{unaryFn = exp, unaryName = "exp", unarySymbol = Nothing})- sqrt :: (Floating a, Columnable a) => Expr a -> Expr a- sqrt =- Unary (MkUnaryOp{unaryFn = sqrt, unaryName = "sqrt", unarySymbol = Nothing})- (**) :: (Floating a, Columnable a) => Expr a -> Expr a -> Expr a- (**) =- Binary- ( MkBinaryOp- { binaryFn = (**)- , binaryName = "exponentiate"- , binarySymbol = Just "**"- , binaryCommutative = False- , binaryPrecedence = 8- }- )- log :: (Floating a, Columnable a) => Expr a -> Expr a- log = Unary (MkUnaryOp{unaryFn = log, unaryName = "log", unarySymbol = Nothing})- logBase :: (Floating a, Columnable a) => Expr a -> Expr a -> Expr a- logBase =- Binary- ( MkBinaryOp- { binaryFn = logBase- , binaryName = "logBase"- , binarySymbol = Nothing- , binaryCommutative = False- , binaryPrecedence = 1- }- )- sin :: (Floating a, Columnable a) => Expr a -> Expr a- sin = Unary (MkUnaryOp{unaryFn = sin, unaryName = "sin", unarySymbol = Nothing})- cos :: (Floating a, Columnable a) => Expr a -> Expr a- cos = Unary (MkUnaryOp{unaryFn = cos, unaryName = "cos", unarySymbol = Nothing})- tan :: (Floating a, Columnable a) => Expr a -> Expr a- tan = Unary (MkUnaryOp{unaryFn = tan, unaryName = "tan", unarySymbol = Nothing})- asin :: (Floating a, Columnable a) => Expr a -> Expr a- asin =- Unary (MkUnaryOp{unaryFn = asin, unaryName = "asin", unarySymbol = Nothing})- acos :: (Floating a, Columnable a) => Expr a -> Expr a- acos =- Unary (MkUnaryOp{unaryFn = acos, unaryName = "acos", unarySymbol = Nothing})- atan :: (Floating a, Columnable a) => Expr a -> Expr a- atan =- Unary (MkUnaryOp{unaryFn = atan, unaryName = "atan", unarySymbol = Nothing})- sinh :: (Floating a, Columnable a) => Expr a -> Expr a- sinh =- Unary (MkUnaryOp{unaryFn = sinh, unaryName = "sinh", unarySymbol = Nothing})- cosh :: (Floating a, Columnable a) => Expr a -> Expr a- cosh =- Unary (MkUnaryOp{unaryFn = cosh, unaryName = "cosh", unarySymbol = Nothing})- asinh :: (Floating a, Columnable a) => Expr a -> Expr a- asinh =- Unary- (MkUnaryOp{unaryFn = asinh, unaryName = "asinh", unarySymbol = Nothing})- acosh :: (Floating a, Columnable a) => Expr a -> Expr a- acosh =- Unary- (MkUnaryOp{unaryFn = acosh, unaryName = "acosh", unarySymbol = Nothing})- atanh :: (Floating a, Columnable a) => Expr a -> Expr a- atanh =- Unary- (MkUnaryOp{unaryFn = atanh, unaryName = "atanh", unarySymbol = Nothing})--instance (Show a) => Show (Expr a) where- show :: Expr a -> String- show (Col name) = "(col @" ++ show (typeRep @a) ++ " " ++ show name ++ ")"- show (CastWith name tag _) = "(castWith " ++ show tag ++ " " ++ show name ++ ")"- show (CastExprWith tag _ inner) = "(castExprWith " ++ show tag ++ " " ++ show inner ++ ")"- show (Lit value) = "(lit (" ++ show value ++ "))"- show (If cond l r) = "(ifThenElse " ++ show cond ++ " " ++ show l ++ " " ++ show r ++ ")"- show (Unary op value) = "(" ++ T.unpack (unaryName op) ++ " " ++ show value ++ ")"- show (Binary op a b) = "(" ++ T.unpack (binaryName op) ++ " " ++ show a ++ " " ++ show b ++ ")"- show (Agg (CollectAgg op _) expr) = "(" ++ T.unpack op ++ " " ++ show expr ++ ")"- show (Agg (FoldAgg op _ _) expr) = "(" ++ T.unpack op ++ " " ++ show expr ++ ")"- show (Agg (MergeAgg op _ _ _ _) expr) = "(" ++ T.unpack op ++ " " ++ show expr ++ ")"- show (Over keys inner) = "(over " ++ show keys ++ " " ++ show inner ++ ")"--normalize :: (Show a, Typeable a) => Expr a -> Expr a-normalize expr = case expr of- Col name -> Col name- CastWith n t f -> CastWith n t f- CastExprWith t f e -> CastExprWith t f (normalize e)- Lit val -> Lit val- If cond th el -> If (normalize cond) (normalize th) (normalize el)- Unary op e -> Unary op (normalize e)- Binary op e1 e2- | binaryCommutative op ->- let n1 = normalize e1- n2 = normalize e2- in case testEquality (typeOf n1) (typeOf n2) of- Nothing -> expr- Just Refl ->- if compareExpr n1 n2 == GT- then Binary op n2 n1 -- Swap to canonical order- else Binary op n1 n2- | otherwise -> Binary op (normalize e1) (normalize e2)- Agg strat e -> Agg strat (normalize e)- Over keys inner -> Over keys (normalize inner)---- Compare expressions for ordering (used in normalization)-compareExpr :: Expr a -> Expr a -> Ordering-compareExpr e1 e2 = compare (exprKey e1) (exprKey e2)- where- exprKey :: Expr a -> String- exprKey (Col name) = "0:" ++ T.unpack name- exprKey (CastWith name tag _) = "0CW:" ++ T.unpack name ++ ":" ++ T.unpack tag- exprKey (CastExprWith tag _ _) = "0CE:" ++ T.unpack tag- exprKey (Lit val) = "1:" ++ show val- exprKey (If c t e) = "2:" ++ exprKey c ++ exprKey t ++ exprKey e- exprKey (Unary op e) = "3:" ++ T.unpack (unaryName op) ++ exprKey e- exprKey (Binary op e1' e2') = "4:" ++ T.unpack (binaryName op) ++ exprKey e1' ++ exprKey e2'- exprKey (Agg (CollectAgg name _) e) = "5:" ++ T.unpack name ++ exprKey e- exprKey (Agg (FoldAgg name _ _) e) = "5:" ++ T.unpack name ++ exprKey e- exprKey (Agg (MergeAgg name _ _ _ _) e) = "5:" ++ T.unpack name ++ exprKey e- exprKey (Over keys e) = "6:over:" ++ show keys ++ exprKey e--eqExpr :: forall a. (Columnable a) => Expr a -> Expr a -> Bool-eqExpr l r = eqNormalized (normalize l) (normalize r)- where- exprEq :: (Columnable b, Columnable c) => Expr b -> Expr c -> Bool- exprEq e1 e2 = case testEquality (typeOf e1) (typeOf e2) of- Just Refl -> eqExpr e1 e2- Nothing -> False- eqNormalized :: Expr a -> Expr a -> Bool- eqNormalized (Col n1) (Col n2) = n1 == n2- eqNormalized (CastWith n1 t1 _) (CastWith n2 t2 _) = n1 == n2 && t1 == t2- eqNormalized (CastExprWith t1 _ e1) (CastExprWith t2 _ e2) = t1 == t2 && e1 `exprEq` e2- eqNormalized (Lit v1) (Lit v2) = v1 == v2- eqNormalized (If c1 t1 e1) (If c2 t2 e2) =- eqExpr c1 c2 && t1 `exprEq` t2 && e1 `exprEq` e2- eqNormalized (Unary op1 e1) (Unary op2 e2) = unaryName op1 == unaryName op2 && e1 `exprEq` e2- eqNormalized (Binary op1 e1a e1b) (Binary op2 e2a e2b) = binaryName op1 == binaryName op2 && e1a `exprEq` e2a && e1b `exprEq` e2b- eqNormalized (Agg (CollectAgg n1 _) e1) (Agg (CollectAgg n2 _) e2) =- n1 == n2 && e1 `exprEq` e2- eqNormalized (Agg (FoldAgg n1 _ _) e1) (Agg (FoldAgg n2 _ _) e2) =- n1 == n2 && e1 `exprEq` e2- eqNormalized (Agg (MergeAgg n1 _ _ _ _) e1) (Agg (MergeAgg n2 _ _ _ _) e2) =- n1 == n2 && e1 `exprEq` e2- eqNormalized (Over k1 e1) (Over k2 e2) = k1 == k2 && e1 `exprEq` e2- eqNormalized _ _ = False--replaceExpr ::- forall a b c.- (Columnable a, Columnable b, Columnable c) =>- Expr a -> Expr b -> Expr c -> Expr c-replaceExpr new old expr = case testEquality (typeRep @b) (typeRep @c) of- Just Refl -> case testEquality (typeRep @a) (typeRep @c) of- Just Refl -> if eqExpr old expr then new else replace'- Nothing -> expr- Nothing -> replace'- where- replace' = case expr of- (Col _) -> expr- (CastWith{}) -> expr- (CastExprWith t f e) -> CastExprWith t f (replaceExpr new old e)- (Lit _) -> expr- (If cond l r) ->- If (replaceExpr new old cond) (replaceExpr new old l) (replaceExpr new old r)- (Unary op value) -> Unary op (replaceExpr new old value)- (Binary op l r) -> Binary op (replaceExpr new old l) (replaceExpr new old r)- (Agg op inner) -> Agg op (replaceExpr new old inner)- (Over keys inner) -> Over keys (replaceExpr new old inner)--eSize :: Expr a -> Int-eSize (Col _) = 1-eSize (CastWith{}) = 1-eSize (CastExprWith _ _ e) = 1 + eSize e-eSize (Lit _) = 1-eSize (If c l r) = 1 + eSize c + eSize l + eSize r-eSize (Unary _ e) = 1 + eSize e-eSize (Binary _ l r) = 1 + eSize l + eSize r-eSize (Agg _strategy expr) = eSize expr + 1-eSize (Over _ inner) = 1 + eSize inner--getColumns :: Expr a -> [T.Text]-getColumns (Col cName) = [cName]-getColumns (CastWith name _ _) = [name]-getColumns (CastExprWith _ _ e) = getColumns e-getColumns _expr@(Lit _) = []-getColumns (If cond l r) = getColumns cond <> getColumns l <> getColumns r-getColumns (Unary _op value) = getColumns value-getColumns (Binary _op l r) = getColumns l <> getColumns r-getColumns (Agg _strategy expr) = getColumns expr-getColumns (Over keys inner) = keys <> getColumns inner--prettyPrint :: Expr a -> String-prettyPrint = go 0 0- where- indent :: Int -> String- indent n = replicate (n * 2) ' '-- go :: Int -> Int -> Expr a -> String- go depth prec expr = case expr of- Col name -> T.unpack name- CastWith name _ _ -> T.unpack name- CastExprWith tag _ inner -> T.unpack tag ++ "(" ++ go depth 0 inner ++ ")"- Lit value -> show value- If cond t e ->- let inner =- "if "- ++ go (depth + 1) 0 cond- ++ "\n"- ++ indent (depth + 1)- ++ "then "- ++ go (depth + 1) 0 t- ++ "\n"- ++ indent (depth + 1)- ++ "else "- ++ go (depth + 1) 0 e- in if prec > 0 then "(" ++ inner ++ ")" else inner- Unary op arg -> case unarySymbol op of- Nothing -> T.unpack (unaryName op) ++ "(" ++ go depth 0 arg ++ ")"- Just sym -> T.unpack sym ++ "(" ++ go depth 0 arg ++ ")"- Binary op l r ->- let p = binaryPrecedence op- inner = case binarySymbol op of- Just name -> go depth p l ++ " " ++ T.unpack name ++ " " ++ go depth p r- Nothing ->- T.unpack (binaryName op) ++ "(" ++ go depth p l ++ ", " ++ go depth p r ++ ")"- in if prec > p then "(" ++ inner ++ ")" else inner- Agg (CollectAgg op _) arg -> T.unpack op ++ "(" ++ go depth 0 arg ++ ")"- Agg (FoldAgg op _ _) arg -> T.unpack op ++ "(" ++ go depth 0 arg ++ ")"- Agg (MergeAgg op _ _ _ _) arg -> T.unpack op ++ "(" ++ go depth 0 arg ++ ")"- Over keys inner -> go depth 0 inner ++ ".over(" ++ show (map T.unpack keys) ++ ")"
− src/DataFrame/Internal/Grouping.hs
@@ -1,192 +0,0 @@-{-# LANGUAGE BangPatterns #-}-{-# LANGUAGE ExplicitNamespaces #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE LambdaCase #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE Strict #-}-{-# LANGUAGE TypeApplications #-}--module DataFrame.Internal.Grouping (- groupBy,- buildRowToGroup,- changingPoints,-) where--import qualified Data.List as L-import qualified Data.Map as M-import qualified Data.Text as T-import qualified Data.Vector as V-import qualified Data.Vector.Algorithms.Radix as VA-import qualified Data.Vector.Unboxed as VU-import qualified Data.Vector.Unboxed.Mutable as VUM--import Control.Exception (throw)-import Control.Monad-import Control.Monad.ST (runST)-import Data.Bits-import Data.Hashable-import Data.Type.Equality (TestEquality (..), type (:~:) (Refl))-import DataFrame.Errors-import DataFrame.Internal.Column (- Column (..),- bitmapTestBit,- )-import DataFrame.Internal.DataFrame (DataFrame (..), GroupedDataFrame (..))-import DataFrame.Internal.Types-import Type.Reflection (typeRep)--{- | O(k * n) groups the dataframe by the given rows aggregating the remaining rows-into vector that should be reduced later.--}-groupBy ::- [T.Text] ->- DataFrame ->- GroupedDataFrame-groupBy names df- | any (`notElem` columnNames df) names =- throw $- ColumnsNotFoundException- (names L.\\ columnNames df)- "groupBy"- (columnNames df)- | nRows df == 0 =- Grouped- df- names- VU.empty- (VU.fromList [0])- VU.empty- | otherwise =- let !vis = VU.map fst valIndices- !os = changingPoints valIndices- !n = nRows df- in Grouped- df- names- vis- os- (buildRowToGroup n vis os)- where- indicesToGroup = M.elems $ M.filterWithKey (\k _ -> k `elem` names) (columnIndices df)- doubleToInt :: Double -> Int- doubleToInt = floor . (* 1000)- valIndices = runST $ do- let n = nRows df- mv <- VUM.new n-- let selectedCols = map (columns df V.!) indicesToGroup-- forM_ selectedCols $ \case- UnboxedColumn _ (v :: VU.Vector a) ->- case testEquality (typeRep @a) (typeRep @Int) of- Just Refl ->- VU.imapM_- ( \i x -> do- (_, !h) <- VUM.unsafeRead mv i- VUM.unsafeWrite mv i (i, hashWithSalt h x)- )- v- Nothing ->- case testEquality (typeRep @a) (typeRep @Double) of- Just Refl ->- VU.imapM_- ( \i d -> do- (_, !h) <- VUM.unsafeRead mv i- VUM.unsafeWrite mv i (i, hashWithSalt h (doubleToInt d))- )- v- Nothing ->- case sIntegral @a of- STrue ->- VU.imapM_- ( \i d -> do- let x :: Int- x = fromIntegral @a @Int d- (_, !h) <- VUM.unsafeRead mv i- VUM.unsafeWrite mv i (i, hashWithSalt h x)- )- v- SFalse ->- case sFloating @a of- STrue ->- VU.imapM_- ( \i d -> do- let x :: Int- x = doubleToInt (realToFrac d :: Double)- (_, !h) <- VUM.unsafeRead mv i- VUM.unsafeWrite mv i (i, hashWithSalt h x)- )- v- SFalse ->- VU.imapM_- ( \i d -> do- let x = hash (show d)- (_, !h) <- VUM.unsafeRead mv i- VUM.unsafeWrite mv i (i, hashWithSalt h x)- )- v- BoxedColumn bm (v :: V.Vector a) ->- case testEquality (typeRep @a) (typeRep @T.Text) of- Just Refl ->- V.imapM_- ( \i t -> do- (_, !h) <- VUM.unsafeRead mv i- let h' = case bm of- Just bm' | not (bitmapTestBit bm' i) -> hashWithSalt h (0 :: Int) -- null sentinel- _ -> hashWithSalt h t- VUM.unsafeWrite mv i (i, h')- )- v- Nothing ->- V.imapM_- ( \i d -> do- (_, !h) <- VUM.unsafeRead mv i- let h' = case bm of- Just bm' | not (bitmapTestBit bm' i) -> hashWithSalt h (0 :: Int) -- null sentinel- _ -> hashWithSalt h (hash (show d))- VUM.unsafeWrite mv i (i, h')- )- v-- let numPasses = 4- bucketSize = 65536- radixFunc k (_, !h) =- let h' = fromIntegral h `xor` (1 `unsafeShiftL` 63) :: Word- shiftBits = k * 16- in fromIntegral ((h' `unsafeShiftR` shiftBits) .&. 65535)- VA.sortBy numPasses bucketSize radixFunc mv- VU.unsafeFreeze mv---- Inline accessors to avoid depending on Operations.Core--columnNames :: DataFrame -> [T.Text]-columnNames = M.keys . columnIndices--nRows :: DataFrame -> Int-nRows = fst . dataframeDimensions--{- | Build the rowToGroup lookup vector from valueIndices and offsets.-rowToGroup[i] = k means row i belongs to group k.--}-buildRowToGroup :: Int -> VU.Vector Int -> VU.Vector Int -> VU.Vector Int-buildRowToGroup n vis os = runST $ do- rtg <- VUM.new n- let nGroups = VU.length os - 1- forM_ [0 .. nGroups - 1] $ \k ->- let s = VU.unsafeIndex os k- e = VU.unsafeIndex os (k + 1)- in forM_ [s .. e - 1] $ \i ->- VUM.unsafeWrite rtg (VU.unsafeIndex vis i) k- VU.unsafeFreeze rtg-{-# NOINLINE buildRowToGroup #-}--changingPoints :: VU.Vector (Int, Int) -> VU.Vector Int-changingPoints vs =- VU.reverse- (VU.fromList (VU.length vs : fst (VU.ifoldl' findChangePoints initialState vs)))- where- initialState = ([0], snd (VU.head vs))- findChangePoints (!offs, !currentVal) index (_, !newVal)- | currentVal == newVal = (offs, currentVal)- | otherwise = (index : offs, newVal)
− src/DataFrame/Internal/Interpreter.hs
@@ -1,1064 +0,0 @@-{-# LANGUAGE AllowAmbiguousTypes #-}-{-# LANGUAGE BangPatterns #-}-{-# LANGUAGE ExplicitNamespaces #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE FlexibleInstances #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE MultiParamTypeClasses #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE RankNTypes #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}-{-# LANGUAGE UndecidableInstances #-}-{-# OPTIONS_GHC -Wno-orphans #-}--module DataFrame.Internal.Interpreter (- -- * New core API- Value (..),- Ctx (..),- eval,- materialize,-- -- * Backward-compatible API- interpret,- interpretAggregation,- AggregationResult (..),-) where--import Data.Bifunctor (first)-import qualified Data.Map as M-import qualified Data.Text as T-import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl))-import qualified Data.Vector as V-import qualified Data.Vector.Generic as VG-import qualified Data.Vector.Unboxed as VU-import qualified Data.Vector.Unboxed.Mutable as VUM-import DataFrame.Errors-import DataFrame.Internal.Column-import DataFrame.Internal.DataFrame-import DataFrame.Internal.Expression-import qualified DataFrame.Internal.Grouping as G-import DataFrame.Internal.Types-import Type.Reflection (- Typeable,- typeRep,- )--import Data.Int (Int16, Int32, Int64, Int8)---- Specializations for common aggregation types to avoid dictionary overhead.--- foldLinearGroups: mean accumulator-{-# SPECIALIZE foldLinearGroups ::- (MeanAcc -> Double -> MeanAcc) ->- MeanAcc ->- Column ->- VU.Vector Int ->- Int ->- Either DataFrameException Column- #-}-{-# SPECIALIZE foldLinearGroups ::- (MeanAcc -> Float -> MeanAcc) ->- MeanAcc ->- Column ->- VU.Vector Int ->- Int ->- Either DataFrameException Column- #-}-{-# SPECIALIZE foldLinearGroups ::- (MeanAcc -> Int -> MeanAcc) ->- MeanAcc ->- Column ->- VU.Vector Int ->- Int ->- Either DataFrameException Column- #-}-{-# SPECIALIZE foldLinearGroups ::- (MeanAcc -> Int8 -> MeanAcc) ->- MeanAcc ->- Column ->- VU.Vector Int ->- Int ->- Either DataFrameException Column- #-}-{-# SPECIALIZE foldLinearGroups ::- (MeanAcc -> Int16 -> MeanAcc) ->- MeanAcc ->- Column ->- VU.Vector Int ->- Int ->- Either DataFrameException Column- #-}-{-# SPECIALIZE foldLinearGroups ::- (MeanAcc -> Int32 -> MeanAcc) ->- MeanAcc ->- Column ->- VU.Vector Int ->- Int ->- Either DataFrameException Column- #-}-{-# SPECIALIZE foldLinearGroups ::- (MeanAcc -> Int64 -> MeanAcc) ->- MeanAcc ->- Column ->- VU.Vector Int ->- Int ->- Either DataFrameException Column- #-}--- foldLinearGroups: count accumulator-{-# SPECIALIZE foldLinearGroups ::- (Int -> Double -> Int) ->- Int ->- Column ->- VU.Vector Int ->- Int ->- Either DataFrameException Column- #-}-{-# SPECIALIZE foldLinearGroups ::- (Int -> Float -> Int) ->- Int ->- Column ->- VU.Vector Int ->- Int ->- Either DataFrameException Column- #-}-{-# SPECIALIZE foldLinearGroups ::- (Int -> Int -> Int) ->- Int ->- Column ->- VU.Vector Int ->- Int ->- Either DataFrameException Column- #-}-{-# SPECIALIZE foldLinearGroups ::- (Int -> Int8 -> Int) ->- Int ->- Column ->- VU.Vector Int ->- Int ->- Either DataFrameException Column- #-}-{-# SPECIALIZE foldLinearGroups ::- (Int -> Int16 -> Int) ->- Int ->- Column ->- VU.Vector Int ->- Int ->- Either DataFrameException Column- #-}-{-# SPECIALIZE foldLinearGroups ::- (Int -> Int32 -> Int) ->- Int ->- Column ->- VU.Vector Int ->- Int ->- Either DataFrameException Column- #-}-{-# SPECIALIZE foldLinearGroups ::- (Int -> Int64 -> Int) ->- Int ->- Column ->- VU.Vector Int ->- Int ->- Either DataFrameException Column- #-}--- foldLinearGroups: sum/min/max (acc == elem)-{-# SPECIALIZE foldLinearGroups ::- (Double -> Double -> Double) ->- Double ->- Column ->- VU.Vector Int ->- Int ->- Either DataFrameException Column- #-}-{-# SPECIALIZE foldLinearGroups ::- (Float -> Float -> Float) ->- Float ->- Column ->- VU.Vector Int ->- Int ->- Either DataFrameException Column- #-}-{-# SPECIALIZE foldLinearGroups ::- (Int8 -> Int8 -> Int8) ->- Int8 ->- Column ->- VU.Vector Int ->- Int ->- Either DataFrameException Column- #-}-{-# SPECIALIZE foldLinearGroups ::- (Int16 -> Int16 -> Int16) ->- Int16 ->- Column ->- VU.Vector Int ->- Int ->- Either DataFrameException Column- #-}-{-# SPECIALIZE foldLinearGroups ::- (Int32 -> Int32 -> Int32) ->- Int32 ->- Column ->- VU.Vector Int ->- Int ->- Either DataFrameException Column- #-}-{-# SPECIALIZE foldLinearGroups ::- (Int64 -> Int64 -> Int64) ->- Int64 ->- Column ->- VU.Vector Int ->- Int ->- Either DataFrameException Column- #-}---- mapColumn: finalize-{-# SPECIALIZE mapColumn ::- (MeanAcc -> Double) -> Column -> Either DataFrameException Column- #-}-{-# SPECIALIZE mapColumn ::- (Double -> Double) -> Column -> Either DataFrameException Column- #-}-{-# SPECIALIZE mapColumn ::- (Float -> Float) -> Column -> Either DataFrameException Column- #-}-{-# SPECIALIZE mapColumn ::- (Int -> Int) -> Column -> Either DataFrameException Column- #-}---- zipWithColumns: binary ops-{-# SPECIALIZE zipWithColumns ::- (Double -> Double -> Double) ->- Column ->- Column ->- Either DataFrameException Column- #-}-{-# SPECIALIZE zipWithColumns ::- (Float -> Float -> Float) ->- Column ->- Column ->- Either DataFrameException Column- #-}-{-# SPECIALIZE zipWithColumns ::- (Int -> Int -> Int) -> Column -> Column -> Either DataFrameException Column- #-}-{-# SPECIALIZE zipWithColumns ::- (Int8 -> Int8 -> Int8) -> Column -> Column -> Either DataFrameException Column- #-}-{-# SPECIALIZE zipWithColumns ::- (Int16 -> Int16 -> Int16) ->- Column ->- Column ->- Either DataFrameException Column- #-}-{-# SPECIALIZE zipWithColumns ::- (Int32 -> Int32 -> Int32) ->- Column ->- Column ->- Either DataFrameException Column- #-}-{-# SPECIALIZE zipWithColumns ::- (Int64 -> Int64 -> Int64) ->- Column ->- Column ->- Either DataFrameException Column- #-}------------------------------------------------------------------------------------ Value: the unified result type----------------------------------------------------------------------------------{- | The result of interpreting an expression. Keeps literals as scalars-until the point where a concrete column is needed, avoiding premature-broadcast allocations.--}-data Value a where- -- | A single value, not yet broadcast to any length.- Scalar :: (Columnable a) => a -> Value a- {- | A flat column (one element per row in the flat case, or one- element per group after aggregation).- -}- Flat :: (Columnable a) => Column -> Value a- {- | A grouped column: one 'Column' slice per group. Only produced- when interpreting inside a 'GroupCtx'.- -}- Group :: (Columnable a) => V.Vector Column -> Value a--instance (Show a) => Show (Value a) where- show (Scalar v) = show v- show (Flat v) = show v- show (Group v) = show v---- | The interpretation context.-data Ctx- = FlatCtx DataFrame- | GroupCtx GroupedDataFrame------------------------------------------------------------------------------------ Materialisation----------------------------------------------------------------------------------{- | Force a 'Value' into a flat 'Column' of the given length. Scalars-are broadcast; flat columns are returned as-is.--}-materialize :: forall a. (Columnable a) => Int -> Value a -> Column-materialize n (Scalar v) = broadcastScalar @a n v-materialize _ (Flat c) = c-materialize _ (Group _) =- error "materialize: cannot flatten a grouped value to a single column"--{- | Replicate a scalar to a column of length @n@, choosing the most-efficient representation.--}-broadcastScalar :: forall a. (Columnable a) => Int -> a -> Column-broadcastScalar n v = case sUnbox @a of- STrue -> fromUnboxedVector (VU.replicate n v)- SFalse -> fromVector (V.replicate n v)------------------------------------------------------------------------------------ Lifting: the core combinators------------------------------------------------------------------------------------ | Apply a pure function to a 'Value'.-liftValue ::- (Columnable b, Columnable a) =>- (b -> a) -> Value b -> Either DataFrameException (Value a)-liftValue f (Scalar v) = Right (Scalar (f v))-liftValue f (Flat col) = Flat <$> mapColumn f col-liftValue f (Group gs) = Group <$> V.mapM (mapColumn f) gs--{- | Apply a binary function to two 'Value's. When one side is a-'Scalar' the operation degenerates to a 'liftValue' — this is how the-old @Binary op (Lit l) right@ special cases are recovered without-explicit pattern matches in the evaluator.--}-liftValue2 ::- (Columnable c, Columnable b, Columnable a) =>- (c -> b -> a) ->- Value c ->- Value b ->- Either DataFrameException (Value a)-liftValue2 f (Scalar l) (Scalar r) = Right (Scalar (f l r))-liftValue2 f (Scalar l) v = liftValue (f l) v-liftValue2 f v (Scalar r) = liftValue (`f` r) v-liftValue2 f (Flat l) (Flat r) = Flat <$> zipWithColumns f l r-liftValue2 f (Group ls) (Group rs)- | V.length ls == V.length rs =- Group <$> V.zipWithM (zipWithColumns f) ls rs--- Shape mismatches: aggregated vs. non-aggregated.-liftValue2 _ (Flat _) (Group _) =- Left $ AggregatedAndNonAggregatedException "aggregated" "non-aggregated"-liftValue2 _ (Group _) (Flat _) =- Left $ AggregatedAndNonAggregatedException "non-aggregated" "aggregated"-liftValue2 _ (Group _) (Group _) =- Left $ InternalException "Group count mismatch in binary operation"---- | Branch on a boolean 'Value', selecting from two same-typed 'Value's.-branchValue ::- forall a.- (Columnable a) =>- Value Bool ->- Value a ->- Value a ->- Either DataFrameException (Value a)-branchValue (Scalar True) l _ = Right l-branchValue (Scalar False) _ r = Right r-branchValue cond (Scalar l) (Scalar r) =- liftValue (\c -> if c then l else r) cond-branchValue cond (Scalar l) r =- liftValue2 (\c rv -> if c then l else rv) cond r-branchValue cond l (Scalar r) =- liftValue2 (\c lv -> if c then lv else r) cond l-branchValue (Flat cc) (Flat lc) (Flat rc) =- Flat <$> branchColumn @a cc lc rc-branchValue (Group cgs) (Group lgs) (Group rgs)- | V.length cgs == V.length lgs- && V.length lgs == V.length rgs =- Group- <$> V.generateM- (V.length cgs)- ( \i ->- branchColumn @a (cgs V.! i) (lgs V.! i) (rgs V.! i)- )-branchValue _ _ _ =- Left $- AggregatedAndNonAggregatedException- "if-then-else branches"- "mismatched shapes"--{- | Low-level column branch: given a boolean column and two same-typed-columns, produce the element-wise selection.--}-branchColumn ::- forall a.- (Columnable a) =>- Column ->- Column ->- Column ->- Either DataFrameException Column-branchColumn cc lc rc = do- cs <- toVector @Bool @V.Vector cc- ls <- toVector @a @V.Vector lc- rs <- toVector @a @V.Vector rc- pure $- fromVector @a $- V.zipWith3 (\c l r -> if c then l else r) cs ls rs------------------------------------------------------------------------------------ Error enrichment----------------------------------------------------------------------------------{- | Wrap an interpretation step so that any 'TypeMismatchException' gets-annotated with the expression that was being evaluated.--}-addContext ::- (Show a) => Expr a -> Either DataFrameException b -> Either DataFrameException b-addContext expr = first (enrichError (show expr))--enrichError :: String -> DataFrameException -> DataFrameException-enrichError loc (TypeMismatchException ctx) =- TypeMismatchException- ctx- { callingFunctionName =- callingFunctionName ctx <|+> Just "eval"- , errorColumnName =- errorColumnName ctx <|+> Just loc- }- where- -- Prefer the existing value; fall back to the new one.- Nothing <|+> b = b- a <|+> _ = a-enrichError _ e = e------------------------------------------------------------------------------------ Group slicing----------------------------------------------------------------------------------{- | Given a flat column and grouping metadata, produce one 'Column' per-group. Each result column is an O(1) slice into a sorted copy of the-input — the sort happens once, not per-group.--}-sliceGroups :: Column -> VU.Vector Int -> VU.Vector Int -> V.Vector Column-sliceGroups col os indices = case col of- BoxedColumn bm vec ->- let !sorted =- V.generate- (VU.length indices)- ((vec `V.unsafeIndex`) . (indices `VU.unsafeIndex`))- in V.generate nGroups $ \i ->- BoxedColumn- (fmap (bitmapSlice (start i) (len i)) bm)- (V.unsafeSlice (start i) (len i) sorted)- UnboxedColumn bm vec ->- let !sorted = VU.unsafeBackpermute vec indices- in V.generate nGroups $ \i ->- UnboxedColumn- (fmap (bitmapSlice (start i) (len i)) bm)- (VU.unsafeSlice (start i) (len i) sorted)- where- !nGroups = VU.length os - 1- start i = os `VU.unsafeIndex` i- len i = os `VU.unsafeIndex` (i + 1) - start i-{-# INLINE sliceGroups #-}--numGroups :: GroupedDataFrame -> Int-numGroups gdf = VU.length (offsets gdf) - 1---- | Build the inverse of a permutation vector.-invertPermutation :: VU.Vector Int -> VU.Vector Int-invertPermutation perm = VU.create $ do- let !n = VU.length perm- inv <- VUM.new n- VU.imapM_ (flip (VUM.unsafeWrite inv)) perm- return inv-{-# INLINE invertPermutation #-}------------------------------------------------------------------------------------ promoteColumnWith: unified numeric / text coercion for CastWith----------------------------------------------------------------------------------{- | Apply a result-handler @onResult@ to each element of a column after-coercing it to type @a@. Covers three modes in one:--* @onResult = either (const Nothing) Just@ → like @cast@ (returns @Maybe a@)-* @onResult = either (const def) id@ → like @castWithDefault@ (returns @a@)-* @onResult = either (Left . T.pack) Right@ → like @castEither@ (returns @Either T.Text a@)--Numeric coercion handles Double, Float, and Int targets. Text columns-(String / T.Text) are parsed via 'reads'. Any other mismatch returns-'Left TypeMismatchException'.--}-promoteColumnWith ::- forall a b.- (Columnable a, Columnable b, Read a) =>- (Either String a -> b) -> Column -> Either DataFrameException Column-promoteColumnWith onResult col- | hasElemType @b col = Right col- | hasElemType @a col = mapColumn @a (onResult . Right) col- | Just result <- tryMaybeWrap @a @b onResult col = result- | otherwise =- case testEquality (typeRep @a) (typeRep @Double) of- Just Refl -> promoteToDoubleWith onResult col- Nothing ->- case testEquality (typeRep @a) (typeRep @Float) of- Just Refl -> promoteToFloatWith onResult col- Nothing ->- case testEquality (typeRep @a) (typeRep @Int) of- Just Refl -> promoteToIntWith onResult col- Nothing -> tryParseWith @a onResult col--promoteToDoubleWith ::- forall b.- (Columnable b) =>- (Either String Double -> b) -> Column -> Either DataFrameException Column-promoteToDoubleWith onResult col = case col of- UnboxedColumn Nothing (v :: VU.Vector c) ->- case sFloating @c of- STrue ->- Right $- fromVector @b- (V.map (onResult . Right . (realToFrac :: c -> Double)) (VG.convert v))- SFalse -> case sIntegral @c of- STrue ->- Right $- fromVector @b- (V.map (onResult . Right . (fromIntegral :: c -> Double)) (VG.convert v))- SFalse -> castMismatch @c @b- UnboxedColumn (Just bm) (v :: VU.Vector c) ->- case sFloating @c of- STrue ->- Right $- fromVector @b- ( V.generate (VU.length v) $ \i ->- if bitmapTestBit bm i- then onResult (Right (realToFrac (VU.unsafeIndex v i) :: Double))- else onResult (Left "null")- )- SFalse -> case sIntegral @c of- STrue ->- Right $- fromVector @b- ( V.generate (VU.length v) $ \i ->- if bitmapTestBit bm i- then onResult (Right (fromIntegral (VU.unsafeIndex v i) :: Double))- else onResult (Left "null")- )- SFalse -> castMismatch @c @b- BoxedColumn _ _ -> tryParseWith @Double onResult col--promoteToFloatWith ::- forall b.- (Columnable b) =>- (Either String Float -> b) -> Column -> Either DataFrameException Column-promoteToFloatWith onResult col = case col of- UnboxedColumn Nothing (v :: VU.Vector c) ->- case sFloating @c of- STrue ->- Right $- fromVector @b- (V.map (onResult . Right . (realToFrac :: c -> Float)) (VG.convert v))- SFalse -> case sIntegral @c of- STrue ->- Right $- fromVector @b- (V.map (onResult . Right . (fromIntegral :: c -> Float)) (VG.convert v))- SFalse -> castMismatch @c @b- UnboxedColumn (Just bm) (v :: VU.Vector c) ->- case sFloating @c of- STrue ->- Right $- fromVector @b- ( V.generate (VU.length v) $ \i ->- if bitmapTestBit bm i- then onResult (Right (realToFrac (VU.unsafeIndex v i) :: Float))- else onResult (Left "null")- )- SFalse -> case sIntegral @c of- STrue ->- Right $- fromVector @b- ( V.generate (VU.length v) $ \i ->- if bitmapTestBit bm i- then onResult (Right (fromIntegral (VU.unsafeIndex v i) :: Float))- else onResult (Left "null")- )- SFalse -> castMismatch @c @b- BoxedColumn _ _ -> tryParseWith @Float onResult col--promoteToIntWith ::- forall b.- (Columnable b) =>- (Either String Int -> b) -> Column -> Either DataFrameException Column-promoteToIntWith onResult col = case col of- UnboxedColumn Nothing (v :: VU.Vector c) ->- case sFloating @c of- STrue ->- Right $- fromVector @b- (V.map (onResult . Right . (round . (realToFrac :: c -> Double))) (VG.convert v))- SFalse -> case sIntegral @c of- STrue ->- Right $- fromVector @b- (V.map (onResult . Right . (fromIntegral :: c -> Int)) (VG.convert v))- SFalse -> castMismatch @c @b- UnboxedColumn (Just bm) (v :: VU.Vector c) ->- case sFloating @c of- STrue ->- Right $- fromVector @b- ( V.generate (VU.length v) $ \i ->- if bitmapTestBit bm i- then onResult (Right (round (realToFrac (VU.unsafeIndex v i) :: Double)))- else onResult (Left "null")- )- SFalse -> case sIntegral @c of- STrue ->- Right $- fromVector @b- ( V.generate (VU.length v) $ \i ->- if bitmapTestBit bm i- then onResult (Right (fromIntegral (VU.unsafeIndex v i) :: Int))- else onResult (Left "null")- )- SFalse -> castMismatch @c @b- BoxedColumn _ _ -> tryParseWith @Int onResult col---- | Single parse primitive: apply @onResult@ to the result of 'reads'.-parseWith :: (Read a) => (Either String a -> b) -> String -> b-parseWith f s = case reads s of- [(x, "")] -> f (Right x)- _ -> case reads (show s) of- [(x, "")] -> f (Right x)- _ -> f (Left s)--tryParseWith ::- forall a b.- (Columnable a, Columnable b, Read a) =>- (Either String a -> b) -> Column -> Either DataFrameException Column-tryParseWith onResult col = case col of- BoxedColumn bm (v :: V.Vector c) ->- case testEquality (typeRep @c) (typeRep @String) of- Just Refl -> case bm of- Nothing -> Right $ fromVector @b $ V.map (parseWith onResult) v- Just bitmap ->- Right $- fromVector @b $- V.imap- ( \i x ->- if bitmapTestBit bitmap i then parseWith onResult x else onResult (Left "null")- )- v- Nothing ->- case testEquality (typeRep @c) (typeRep @T.Text) of- Just Refl -> case bm of- Nothing -> Right $ fromVector @b $ V.map (parseWith onResult . T.unpack) v- Just bitmap ->- Right $- fromVector @b $- V.imap- ( \i x ->- if bitmapTestBit bitmap i- then parseWith onResult (T.unpack x)- else onResult (Left "null")- )- v- Nothing -> castMismatch @c @b- UnboxedColumn bm (v :: VU.Vector c) -> case bm of- Nothing -> Right $ fromVector @b $ V.map (parseWith onResult . show) (V.convert v)- Just bitmap ->- Right $- fromVector @b $- V.imap- ( \i x ->- if bitmapTestBit bitmap i- then parseWith onResult (show x)- else onResult (Left "null")- )- (V.convert v)--{- | When the output type @b@ is @Maybe c@ (or @Maybe (Maybe c)@) and the-column stores plain @c@ values, wrap each element in 'Just'.-The @Maybe (Maybe c)@ case applies join semantics: instead of producing-a double-wrapped column, a @Maybe c@ column is returned, so-@castExpr \@(Maybe Double)@ on a @Double@ column yields @Maybe Double@-rather than @Maybe (Maybe Double)@.-Returns 'Nothing' when neither condition holds.--}-tryMaybeWrap ::- forall a b.- (Columnable a, Columnable b) =>- (Either String a -> b) -> Column -> Maybe (Either DataFrameException Column)-tryMaybeWrap _onResult col = case col of- UnboxedColumn Nothing (v :: VU.Vector c) ->- let wrapped = V.map Just (VG.convert v) :: V.Vector (Maybe c)- in case testEquality (typeRep @b) (typeRep @(Maybe c)) of- Just Refl -> Just $ Right $ fromVector @b wrapped- Nothing ->- case testEquality (typeRep @b) (typeRep @(Maybe (Maybe c))) of- Just _ -> Just $ Right $ fromVector @(Maybe c) wrapped- Nothing -> Nothing- BoxedColumn Nothing (v :: V.Vector c) ->- let wrapped = V.map Just v :: V.Vector (Maybe c)- in case testEquality (typeRep @b) (typeRep @(Maybe c)) of- Just Refl -> Just $ Right $ fromVector @b wrapped- Nothing ->- case testEquality (typeRep @b) (typeRep @(Maybe (Maybe c))) of- Just _ -> Just $ Right $ fromVector @(Maybe c) wrapped- Nothing -> Nothing- _ -> Nothing--castMismatch ::- forall src tgt.- (Typeable src, Typeable tgt) =>- Either DataFrameException Column-castMismatch =- Left $- TypeMismatchException- MkTypeErrorContext- { userType = Right (typeRep @tgt)- , expectedType = Right (typeRep @src)- , callingFunctionName = Just "cast"- , errorColumnName = Nothing- }------------------------------------------------------------------------------------ eval: the unified interpreter----------------------------------------------------------------------------------{- | Evaluate an expression in a given context, producing a 'Value'.-This single function replaces both the old @interpret@ (flat) and-@interpretAggregation@ (grouped) code paths.--}-eval ::- forall a.- (Columnable a) =>- Ctx -> Expr a -> Either DataFrameException (Value a)--- Leaves -------------------------------------------------------------------eval _ (Lit v) = Right (Scalar v)-eval (FlatCtx df) (Col name) =- case getColumn name df of- Nothing ->- Left $ ColumnsNotFoundException [name] "" (M.keys $ columnIndices df)- Just c- | hasElemType @a c -> Right (Flat c)- | otherwise ->- Left $- TypeMismatchException- ( MkTypeErrorContext- { userType = Right (typeRep @a)- , expectedType = Left (columnTypeString c)- , errorColumnName = Just (T.unpack name)- , callingFunctionName = Just "col"- } ::- TypeErrorContext a ()- )-eval (GroupCtx gdf) (Col name) =- case getColumn name (fullDataframe gdf) of- Nothing ->- Left $- ColumnsNotFoundException- [name]- ""- (M.keys $ columnIndices $ fullDataframe gdf)- Just c- | hasElemType @a c ->- Right (Group (sliceGroups c (offsets gdf) (valueIndices gdf)))- | otherwise ->- Left $- TypeMismatchException- ( MkTypeErrorContext- { userType = Right (typeRep @a)- , expectedType = Left (columnTypeString c)- , errorColumnName = Just (T.unpack name)- , callingFunctionName = Just "col"- } ::- TypeErrorContext a ()- )--- CastWith -----------------------------------------------------------------eval (FlatCtx df) (CastWith name _tag onResult) =- case getColumn name df of- Nothing ->- Left $- ColumnsNotFoundException [name] "" (M.keys $ columnIndices df)- Just c -> Flat <$> promoteColumnWith onResult c-eval (GroupCtx gdf) (CastWith name _tag onResult) =- case getColumn name (fullDataframe gdf) of- Nothing ->- Left $- ColumnsNotFoundException- [name]- ""- (M.keys $ columnIndices $ fullDataframe gdf)- Just c -> do- promoted <- promoteColumnWith onResult c- Right $ Group (sliceGroups promoted (offsets gdf) (valueIndices gdf))--- CastExprWith -------------------------------------------------------------eval ctx (CastExprWith _tag onResult (inner :: Expr src)) = do- v <- eval @src ctx inner- case v of- Scalar s ->- Flat <$> promoteColumnWith onResult (fromList @src [s])- Flat col ->- Flat <$> promoteColumnWith onResult col- Group gs ->- Group <$> V.mapM (promoteColumnWith onResult) gs--- Unary --------------------------------------------------------------------eval ctx expr@(Unary (op :: UnaryOp b a) inner) = addContext expr $ do- v <- eval @b ctx inner- liftValue (unaryFn op) v---- Binary -------------------------------------------------------------------eval ctx expr@(Binary (op :: BinaryOp c b a) left right) =- addContext expr $ do- l <- eval @c ctx left- r <- eval @b ctx right- liftValue2 (binaryFn op) l r---- If -----------------------------------------------------------------------eval ctx expr@(If cond l r) = addContext expr $ do- c <- eval @Bool ctx cond- lv <- eval @a ctx l- rv <- eval @a ctx r- branchValue c lv rv---- Over (window function) ---------------------------------------------------eval (FlatCtx df) expr@(Over keys inner) = addContext expr $ do- let gdf = G.groupBy keys df- v <- eval (GroupCtx gdf) inner- case v of- Scalar s ->- Right (Scalar s)- Flat groupCol ->- -- Scalar agg (mean, sum, median): one value per group.- -- Broadcast via rowToGroup: row i gets value at group rowToGroup[i].- Right (Flat (atIndicesStable (rowToGroup gdf) groupCol))- Group groupCols -> do- -- Concatenate in sorted order, then unsort to original row order.- sorted <- V.fold1M' concatColumns groupCols- let inv = invertPermutation (valueIndices gdf)- Right (Flat (atIndicesStable inv sorted))-eval (GroupCtx _) expr@(Over _ _) =- addContext expr $- Left- ( InternalException- "Over (window function) is not supported inside a grouped context"- )--- Fast path: FoldAgg (seeded) on a bare Col in GroupCtx.--- Avoids the O(n) backpermute in sliceGroups by folding directly over--- permuted indices. Only matches when inner is exactly (Col name).--eval (GroupCtx gdf) expr@(Agg (FoldAgg _ (Just seed) (f :: a -> b -> a)) (Col name :: Expr b)) =- addContext expr $- case getColumn name (fullDataframe gdf) of- Nothing ->- Left $- ColumnsNotFoundException- [name]- ""- (M.keys $ columnIndices $ fullDataframe gdf)- Just col ->- Flat <$> foldLinearGroups @b @a f seed col (rowToGroup gdf) (numGroups gdf)--- Fast path: FoldAgg (seedless) on a bare Col in GroupCtx.--eval (GroupCtx gdf) expr@(Agg (FoldAgg _ Nothing (f :: a -> b -> a)) (Col name :: Expr b)) =- addContext expr $- case testEquality (typeRep @a) (typeRep @b) of- Nothing ->- Left $- InternalException- "Type mismatch in seedless fold: \- \accumulator and element types must match"- Just Refl ->- case getColumn name (fullDataframe gdf) of- Nothing ->- Left $- ColumnsNotFoundException- [name]- ""- (M.keys $ columnIndices $ fullDataframe gdf)- Just col ->- Flat <$> foldl1DirectGroups @b f col (valueIndices gdf) (offsets gdf)--- Fast path: MergeAgg on a bare Col in GroupCtx.--eval- (GroupCtx gdf)- expr@( Agg- (MergeAgg _ seed (step :: acc -> b -> acc) _ (finalize :: acc -> a))- (Col name :: Expr b)- ) =- addContext expr $- case getColumn name (fullDataframe gdf) of- Nothing ->- Left $- ColumnsNotFoundException- [name]- ""- (M.keys $ columnIndices $ fullDataframe gdf)- Just col ->- Flat- <$> ( foldLinearGroups @b step seed col (rowToGroup gdf) (numGroups gdf)- >>= mapColumn finalize- )--- Aggregation: CollectAgg --------------------------------------------------eval ctx expr@(Agg (CollectAgg _ (f :: v b -> a)) inner) =- addContext expr $ do- v <- eval @b ctx inner- case v of- Scalar _ ->- Left $- InternalException- "Cannot apply a collection aggregation to a scalar"- Flat col ->- Scalar <$> applyCollect @v @b @a f col- Group gs ->- Flat . fromVector- <$> V.mapM (applyCollect @v @b @a f) gs---- Aggregation: FoldAgg with seed -------------------------------------------eval ctx expr@(Agg (FoldAgg _ (Just seed) (f :: a -> b -> a)) inner) =- addContext expr $ do- v <- eval @b ctx inner- case v of- Scalar x -> Right (broadcastFold ctx seed f x)- Flat col ->- Scalar <$> foldlColumn @b @a f seed col- Group gs ->- Flat . fromVector- <$> V.mapM (foldlColumn @b @a f seed) gs---- Aggregation: MergeAgg ----------------------------------------------------eval- ctx- expr@( Agg- (MergeAgg _ seed (step :: acc -> b -> acc) _ (finalize :: acc -> a))- (inner :: Expr b)- ) =- addContext expr $ do- v <- eval @b ctx inner- case v of- Scalar x -> case broadcastFold ctx seed step x of- Scalar acc -> Right (Scalar (finalize acc))- Flat col -> Flat <$> mapColumn @acc @a finalize col- Group _ ->- Left- ( InternalException- "broadcastFold unexpectedly produced a Group value"- )- Flat col ->- Scalar . finalize <$> foldlColumn @b step seed col- Group gs ->- Flat . fromVector- <$> V.mapM (fmap finalize . foldlColumn @b step seed) gs---- Aggregation: FoldAgg without seed (fold1) --------------------------------eval ctx expr@(Agg (FoldAgg _ Nothing (f :: a -> b -> a)) inner) =- addContext expr $- case testEquality (typeRep @a) (typeRep @b) of- Nothing ->- Left $- InternalException- "Type mismatch in seedless fold: \- \accumulator and element types must match"- Just Refl -> do- v <- eval @b ctx inner- case v of- Scalar _ ->- Left $- InternalException- "fold1 requires at least one element"- Flat col ->- Scalar <$> foldl1Column @a f col- Group gs ->- Flat . fromVector- <$> V.mapM (foldl1Column @a f) gs--broadcastFold ::- forall acc b.- (Columnable acc) =>- Ctx -> acc -> (acc -> b -> acc) -> b -> Value acc-broadcastFold (FlatCtx df) seed step x =- let n = fst (dataframeDimensions df)- in Scalar (iterateStep n step seed x)-broadcastFold (GroupCtx gdf) seed step x =- let offs = offsets gdf- ng = VU.length offs - 1- results =- V.generate ng $ \i ->- let sz = offs VU.! (i + 1) - offs VU.! i- in iterateStep sz step seed x- in Flat (fromVector results)--iterateStep :: Int -> (acc -> b -> acc) -> acc -> b -> acc-iterateStep n step = go n- where- go 0 !acc _ = acc- go k !acc x = go (k - 1) (step acc x) x--{- | Apply a 'CollectAgg' function to a single column, extracting the-appropriate vector type and applying the aggregation function.--}-applyCollect ::- forall v b a.- (VG.Vector v b, Typeable v, Columnable b, Columnable a) =>- (v b -> a) -> Column -> Either DataFrameException a-applyCollect f col = f <$> toVector @b @v col--{- | Result of interpreting an expression in a grouped context.-Retained for backward compatibility with 'aggregate' and friends.--}-data AggregationResult a- = UnAggregated Column- | Aggregated (TypedColumn a)--{- | Interpret an expression against a flat 'DataFrame', producing a-typed column. This is the original top-level entry point; internally-it calls 'eval' and materialises the result.--NOTE: unlike the old implementation, 'Lit' values are no longer-eagerly broadcast. The broadcast happens here, at the boundary,-via 'materialize'.--}-interpret ::- forall a.- (Columnable a) =>- DataFrame -> Expr a -> Either DataFrameException (TypedColumn a)-interpret df expr = do- v <- eval (FlatCtx df) expr- pure $ TColumn $ materialize @a (fst (dataframeDimensions df)) v--{- | Interpret an expression against a 'GroupedDataFrame',-distinguishing aggregated results from bare column references.-Internally calls 'eval'.--}-interpretAggregation ::- forall a.- (Columnable a) =>- GroupedDataFrame ->- Expr a ->- Either DataFrameException (AggregationResult a)-interpretAggregation gdf expr = do- v <- eval (GroupCtx gdf) expr- case v of- Scalar a ->- Right $- Aggregated $- TColumn $- broadcastScalar @a (numGroups gdf) a- Flat col ->- Right $ Aggregated $ TColumn col- Group _ ->- -- The Column payload is intentionally unused — the only- -- call-site ('aggregate') immediately throws- -- 'UnaggregatedException' on this constructor.- Right $ UnAggregated $ BoxedColumn @T.Text Nothing V.empty
− src/DataFrame/Internal/Nullable.hs
@@ -1,500 +0,0 @@-{-# LANGUAGE AllowAmbiguousTypes #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE FlexibleInstances #-}-{-# LANGUAGE FunctionalDependencies #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}-{-# LANGUAGE TypeFamilies #-}-{-# LANGUAGE TypeOperators #-}-{-# LANGUAGE UndecidableInstances #-}-{-# LANGUAGE UndecidableSuperClasses #-}--{- | Nullable-aware binary operations for expressions.--This module provides two type classes, 'NullableArithOp' and 'NullableCmpOp',-which enable operators like '.+', '.-', '.*', './', '.==' etc. to work-transparently across combinations of nullable (@Maybe a@) and non-nullable-(@a@) column types.--The partial functional dependencies uniquely determine the result type from-the operand types, so GHC infers it without annotations.--The four combinations covered for each class:--* @(a, a)@ — non-nullable × non-nullable-* @(Maybe a, a)@ — nullable × non-nullable-* @(a, Maybe a)@ — non-nullable × nullable-* @(Maybe a, Maybe a)@ — both nullable--== Usage--@--- Mixing nullable and non-nullable columns:-F.col \@Int \"x\" '.+' F.col \@(Maybe Int) \"y\" -- :: Expr (Maybe Int)---- Both non-nullable (existing behaviour preserved):-F.col \@Int \"x\" '.+' F.col \@Int \"y\" -- :: Expr Int---- Comparison with three-valued logic:-F.col \@(Maybe Int) \"x\" '.==' F.col \@Int \"y\" -- :: Expr (Maybe Bool)-@--}-module DataFrame.Internal.Nullable (- -- * Type family- BaseType,-- -- * Arithmetic class- NullableArithOp (..),-- -- * Comparison class- NullableCmpOp (..),-- -- * Generalized nullable lift classes- NullLift1Op (..),- NullLift2Op (..),-- -- * Result-type type families (drive inference in nullLift / nullLift2)- NullLift1Result,- NullLift2Result,-- -- * Result-type type family for comparison operators- NullCmpResult,-- -- * Numeric widening- NumericWidenOp (..),- widenArithOp,- widenCmpOp,- WidenResult,-- -- * Division widening (integral × integral → Double)- DivWidenOp (..),- divArithOp,- WidenResultDiv,-) where--import Data.Int (Int32, Int64)-import DataFrame.Internal.Column (Columnable)-import DataFrame.Internal.Types (Promote, PromoteDiv)--{- | Strip one layer of 'Maybe'.--@-BaseType (Maybe a) = a-BaseType a = a -- for any non-Maybe type-@--}-type family BaseType a where- BaseType (Maybe a) = a- BaseType a = a--{- | Class for arithmetic binary operations that work transparently over-nullable and non-nullable column types.--The functional dependency @a b -> c@ ensures GHC can infer the result type @c@-from the operand types. The 'OVERLAPPABLE' pragma on the non-nullable instance-ensures the more specific @(Maybe a, Maybe a)@ instance wins when both operands-are nullable.--}-class- ( Columnable a- , Columnable b- , Columnable c- ) =>- NullableArithOp a b c- | a b -> c- where- {- | Lift an arithmetic function over the inner (non-Maybe) values.- 'Nothing' short-circuits: any 'Nothing' operand produces 'Nothing'.- -}- nullArithOp ::- (BaseType a -> BaseType a -> BaseType a) ->- a ->- b ->- c--{- | Compute the result type of a nullable comparison.--@-NullCmpResult (Maybe a) b = Maybe Bool-NullCmpResult a (Maybe b) = Maybe Bool -- when a is apart from Maybe-NullCmpResult a b = Bool-@--Used by the comparison operators ('.==', '.<', etc.) so GHC infers the-return type without an explicit annotation.--}-type family NullCmpResult a b where- NullCmpResult (Maybe a) b = Maybe Bool- NullCmpResult a (Maybe b) = Maybe Bool- NullCmpResult a b = Bool--{- | Class for comparison binary operations that work transparently over-nullable and non-nullable column types.--No functional dependency on @e@: the 'OVERLAPPING'\/'OVERLAPPABLE' pragmas on-instances disambiguate at call sites without a FundDep (which would conflict-when both operands are @Maybe@). GHC selects the unique most-specific instance-from the concrete operand types.--}-class- ( Columnable a- , Columnable b- , Columnable e- ) =>- NullableCmpOp a b e- where- {- | Lift a comparison function over the inner values (three-valued logic).- Returns 'Nothing' when either operand is 'Nothing'.- -}- nullCmpOp ::- (BaseType a -> BaseType a -> Bool) ->- a ->- b ->- e--{- | Non-nullable × Non-nullable: apply directly, no wrapping.-Arithmetic result is @a@; comparison result is @Bool@.--}-instance- {-# OVERLAPPABLE #-}- (Columnable a, a ~ BaseType a) =>- NullableArithOp a a a- where- nullArithOp f = f--instance- {-# OVERLAPPABLE #-}- (Columnable a, Columnable Bool, a ~ BaseType a) =>- NullableCmpOp a a Bool- where- nullCmpOp f = f---- | Nullable × Non-nullable: 'Nothing' short-circuits.-instance- (Columnable a, Columnable (Maybe a)) =>- NullableArithOp (Maybe a) a (Maybe a)- where- nullArithOp _f Nothing _ = Nothing- nullArithOp f (Just x) y = Just (f x y)--instance- (Columnable a, Columnable (Maybe a), Columnable (Maybe Bool)) =>- NullableCmpOp (Maybe a) a (Maybe Bool)- where- nullCmpOp _f Nothing _ = Nothing- nullCmpOp f (Just x) y = Just (f x y)---- | Non-nullable × Nullable: 'Nothing' short-circuits.-instance- ( Columnable a- , Columnable (Maybe a)- , a ~ BaseType a- ) =>- NullableArithOp a (Maybe a) (Maybe a)- where- nullArithOp _f _ Nothing = Nothing- nullArithOp f x (Just y) = Just (f x y)--instance- ( Columnable a- , Columnable (Maybe a)- , Columnable (Maybe Bool)- , a ~ BaseType a- ) =>- NullableCmpOp a (Maybe a) (Maybe Bool)- where- nullCmpOp _f _ Nothing = Nothing- nullCmpOp f x (Just y) = Just (f x y)---- | Nullable × Nullable: either 'Nothing' short-circuits.-instance- {-# OVERLAPPING #-}- (Columnable a, Columnable (Maybe a)) =>- NullableArithOp (Maybe a) (Maybe a) (Maybe a)- where- nullArithOp _f Nothing _ = Nothing- nullArithOp _f _ Nothing = Nothing- nullArithOp f (Just x) (Just y) = Just (f x y)--instance- {-# OVERLAPPING #-}- (Columnable a, Columnable (Maybe a), Columnable (Maybe Bool)) =>- NullableCmpOp (Maybe a) (Maybe a) (Maybe Bool)- where- nullCmpOp _f Nothing _ = Nothing- nullCmpOp _f _ Nothing = Nothing- nullCmpOp f (Just x) (Just y) = Just (f x y)---- ------------------------------------------------------------------------------ Generalized nullable lift (unary)--- -----------------------------------------------------------------------------{- | Lift a unary function over a column expression, propagating 'Nothing'.--When @a@ is non-nullable the function is applied directly; when @a = Maybe x@-the function is applied under the 'Just' and 'Nothing' short-circuits.--Use via 'DataFrame.Functions.nullLift'.--}--{- | Compute the result type of a nullable unary lift.--@-NullLift1Result (Maybe a) r = Maybe r-NullLift1Result a r = r -- for any non-Maybe a-@--Used by 'DataFrame.Functions.nullLift' so GHC can infer the return type-without an explicit annotation.--}-type family NullLift1Result a r where- NullLift1Result (Maybe a) r = Maybe r- NullLift1Result a r = r--class- ( Columnable a- , Columnable r- , Columnable c- ) =>- NullLift1Op a r c- where- applyNull1 :: (BaseType a -> r) -> a -> c---- | Non-nullable: apply directly.-instance- {-# OVERLAPPABLE #-}- (Columnable a, Columnable r, a ~ BaseType a) =>- NullLift1Op a r r- where- applyNull1 f = f---- | Nullable: propagate 'Nothing'.-instance- {-# OVERLAPPING #-}- (Columnable a, Columnable r, Columnable (Maybe r)) =>- NullLift1Op (Maybe a) r (Maybe r)- where- applyNull1 _ Nothing = Nothing- applyNull1 f (Just x) = Just (f x)---- ------------------------------------------------------------------------------ Generalized nullable lift (binary)--- -----------------------------------------------------------------------------{- | Lift a binary function over two column expressions, propagating 'Nothing'.--The four combinations:--* @(a, b)@ — both non-nullable: result is @r@-* @(Maybe a, b)@ — left nullable: result is @Maybe r@-* @(a, Maybe b)@ — right nullable: result is @Maybe r@-* @(Maybe a, Maybe b)@ — both nullable: result is @Maybe r@--Use via 'DataFrame.Functions.nullLift2'.--}--{- | Compute the result type of a nullable binary lift.--@-NullLift2Result (Maybe a) b r = Maybe r-NullLift2Result a (Maybe b) r = Maybe r -- when a is apart from Maybe-NullLift2Result a b r = r-@--Used by 'DataFrame.Functions.nullLift2' so GHC can infer the return type.--}-type family NullLift2Result a b r where- NullLift2Result (Maybe a) b r = Maybe r- NullLift2Result a (Maybe b) r = Maybe r- NullLift2Result a b r = r--class- ( Columnable a- , Columnable b- , Columnable r- , Columnable c- ) =>- NullLift2Op a b r c- where- applyNull2 :: (BaseType a -> BaseType b -> r) -> a -> b -> c---- | Both non-nullable: apply directly.-instance- {-# OVERLAPPABLE #-}- (Columnable a, Columnable b, Columnable r, a ~ BaseType a, b ~ BaseType b) =>- NullLift2Op a b r r- where- applyNull2 f = f---- | Left nullable: 'Nothing' short-circuits.-instance- {-# OVERLAPPABLE #-}- (Columnable a, Columnable b, Columnable r, Columnable (Maybe r), b ~ BaseType b) =>- NullLift2Op (Maybe a) b r (Maybe r)- where- applyNull2 _ Nothing _ = Nothing- applyNull2 f (Just x) y = Just (f x y)---- | Right nullable: 'Nothing' short-circuits.-instance- {-# OVERLAPPABLE #-}- (Columnable a, Columnable b, Columnable r, Columnable (Maybe r), a ~ BaseType a) =>- NullLift2Op a (Maybe b) r (Maybe r)- where- applyNull2 _ _ Nothing = Nothing- applyNull2 f x (Just y) = Just (f x y)---- | Both nullable: either 'Nothing' short-circuits.-instance- {-# OVERLAPPING #-}- (Columnable a, Columnable b, Columnable r, Columnable (Maybe r)) =>- NullLift2Op (Maybe a) (Maybe b) r (Maybe r)- where- applyNull2 _ Nothing _ = Nothing- applyNull2 _ _ Nothing = Nothing- applyNull2 f (Just x) (Just y) = Just (f x y)---- ------------------------------------------------------------------------------ Numeric widening--- -----------------------------------------------------------------------------{- | Widen two numeric base types to their promoted common type.--When @a ~ b@ the coercions are identity; otherwise one operand is widened-(e.g. 'Int' → 'Double').--}-class (Columnable (Promote a b)) => NumericWidenOp a b where- widen1 :: a -> Promote a b- widen2 :: b -> Promote a b---- | Same type: identity coercions.-instance {-# OVERLAPPING #-} (Columnable a) => NumericWidenOp a a where- widen1 = id- widen2 = id--instance NumericWidenOp Int Double where widen1 = fromIntegral; widen2 = id-instance NumericWidenOp Double Int where- widen1 = id- widen2 = fromIntegral-instance NumericWidenOp Float Double where widen1 = realToFrac; widen2 = id-instance NumericWidenOp Double Float where- widen1 = id- widen2 = realToFrac-instance NumericWidenOp Int32 Float where widen1 = fromIntegral; widen2 = id-instance NumericWidenOp Float Int32 where- widen1 = id- widen2 = fromIntegral-instance NumericWidenOp Int32 Double where widen1 = fromIntegral; widen2 = id-instance NumericWidenOp Double Int32 where- widen1 = id- widen2 = fromIntegral-instance NumericWidenOp Int64 Float where widen1 = fromIntegral; widen2 = id-instance NumericWidenOp Float Int64 where- widen1 = id- widen2 = fromIntegral-instance NumericWidenOp Int64 Double where widen1 = fromIntegral; widen2 = id-instance NumericWidenOp Double Int64 where- widen1 = id- widen2 = fromIntegral---- | Apply an arithmetic function after widening both operands to their common type.-widenArithOp ::- forall a b.- (NumericWidenOp a b) =>- (Promote a b -> Promote a b -> Promote a b) ->- a ->- b ->- Promote a b-widenArithOp f x y = f (widen1 @a @b x) (widen2 @a @b y)---- | Apply a comparison function after widening both operands to their common type.-widenCmpOp ::- forall a b.- (NumericWidenOp a b) =>- (Promote a b -> Promote a b -> Bool) ->- a ->- b ->- Bool-widenCmpOp f x y = f (widen1 @a @b x) (widen2 @a @b y)---- | Result type of a widening binary operator, accounting for nullable wrappers.-type WidenResult a b = NullLift2Result a b (Promote (BaseType a) (BaseType b))---- ------------------------------------------------------------------------------ Division widening (integral × integral → Double)--- -----------------------------------------------------------------------------{- | Like 'NumericWidenOp' but uses 'PromoteDiv': integral×integral → Double.-Floating types still dominate (Double > Float), and any two integral types-(same or mixed) are both widened to Double.--}-class (Columnable (PromoteDiv a b)) => DivWidenOp a b where- divWiden1 :: a -> PromoteDiv a b- divWiden2 :: b -> PromoteDiv a b---- Floating same-type (identity)-instance DivWidenOp Double Double where divWiden1 = id; divWiden2 = id-instance DivWidenOp Float Float where divWiden1 = id; divWiden2 = id---- Mixed Double/Float-instance DivWidenOp Double Float where divWiden1 = id; divWiden2 = realToFrac-instance DivWidenOp Float Double where divWiden1 = realToFrac; divWiden2 = id---- Double beats integral-instance DivWidenOp Double Int where divWiden1 = id; divWiden2 = fromIntegral-instance DivWidenOp Int Double where divWiden1 = fromIntegral; divWiden2 = id-instance DivWidenOp Double Int32 where divWiden1 = id; divWiden2 = fromIntegral-instance DivWidenOp Int32 Double where divWiden1 = fromIntegral; divWiden2 = id-instance DivWidenOp Double Int64 where divWiden1 = id; divWiden2 = fromIntegral-instance DivWidenOp Int64 Double where divWiden1 = fromIntegral; divWiden2 = id---- Float beats integral-instance DivWidenOp Float Int where divWiden1 = id; divWiden2 = fromIntegral-instance DivWidenOp Int Float where divWiden1 = fromIntegral; divWiden2 = id-instance DivWidenOp Float Int32 where divWiden1 = id; divWiden2 = fromIntegral-instance DivWidenOp Int32 Float where divWiden1 = fromIntegral; divWiden2 = id-instance DivWidenOp Float Int64 where divWiden1 = id; divWiden2 = fromIntegral-instance DivWidenOp Int64 Float where divWiden1 = fromIntegral; divWiden2 = id---- Integral × integral → Double-instance DivWidenOp Int Int where- divWiden1 = fromIntegral- divWiden2 = fromIntegral-instance DivWidenOp Int32 Int32 where- divWiden1 = fromIntegral- divWiden2 = fromIntegral-instance DivWidenOp Int64 Int64 where- divWiden1 = fromIntegral- divWiden2 = fromIntegral-instance DivWidenOp Int Int32 where- divWiden1 = fromIntegral- divWiden2 = fromIntegral-instance DivWidenOp Int32 Int where- divWiden1 = fromIntegral- divWiden2 = fromIntegral-instance DivWidenOp Int Int64 where- divWiden1 = fromIntegral- divWiden2 = fromIntegral-instance DivWidenOp Int64 Int where- divWiden1 = fromIntegral- divWiden2 = fromIntegral-instance DivWidenOp Int32 Int64 where- divWiden1 = fromIntegral- divWiden2 = fromIntegral-instance DivWidenOp Int64 Int32 where- divWiden1 = fromIntegral- divWiden2 = fromIntegral---- | Apply an arithmetic function after widening both operands via 'PromoteDiv'.-divArithOp ::- forall a b.- (DivWidenOp a b) =>- (PromoteDiv a b -> PromoteDiv a b -> PromoteDiv a b) ->- a ->- b ->- PromoteDiv a b-divArithOp f x y = f (divWiden1 @a @b x) (divWiden2 @a @b y)---- | Result type of a division-widening binary operator, accounting for nullable wrappers.-type WidenResultDiv a b =- NullLift2Result a b (PromoteDiv (BaseType a) (BaseType b))
− src/DataFrame/Internal/Parsing.hs
@@ -1,219 +0,0 @@-{-# LANGUAGE LambdaCase #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE ScopedTypeVariables #-}--module DataFrame.Internal.Parsing where--import qualified Data.ByteString.Char8 as C-import qualified Data.Set as S-import qualified Data.Text as T-import qualified Data.Text.IO as TIO--import Control.Applicative (many, (<|>))-import Data.Attoparsec.Text hiding (decimal, double, signed)-import Data.ByteString.Lex.Fractional-import Data.Foldable (fold)-import Data.Text.Read (decimal, double, signed)-import Data.Time (Day, defaultTimeLocale, parseTimeM)-import GHC.Stack (HasCallStack)-import System.IO (Handle, IOMode (..), hIsEOF, hTell, withFile)-import Prelude hiding (takeWhile)--isNullish :: T.Text -> Bool-isNullish =- ( `S.member`- S.fromList- ["Nothing", "NULL", "", " ", "nan", "null", "N/A", "NaN", "NAN", "NA"]- )--isNullishBS :: C.ByteString -> Bool-isNullishBS =- ( `S.member`- S.fromList- ["Nothing", "NULL", "", " ", "nan", "null", "N/A", "NaN", "NAN", "NA"]- )--isTrueish :: T.Text -> Bool-isTrueish t = t `elem` ["True", "true", "TRUE"]--isFalseish :: T.Text -> Bool-isFalseish t = t `elem` ["False", "false", "FALSE"]--readBool :: (HasCallStack) => T.Text -> Maybe Bool-readBool s- | isTrueish s = Just True- | isFalseish s = Just False- | otherwise = Nothing--readByteStringBool :: C.ByteString -> Maybe Bool-readByteStringBool s- | s `elem` ["True", "true", "TRUE"] = Just True- | s `elem` ["False", "false", "FALSE"] = Just False- | otherwise = Nothing--readByteStringDate :: String -> C.ByteString -> Maybe Day-readByteStringDate fmt = parseTimeM True defaultTimeLocale fmt . C.unpack--readInteger :: (HasCallStack) => T.Text -> Maybe Integer-readInteger s = case signed decimal (T.strip s) of- Left _ -> Nothing- Right (value, "") -> Just value- Right (_value, _) -> Nothing--readInt :: (HasCallStack) => T.Text -> Maybe Int-readInt s = case signed decimal (T.strip s) of- Left _ -> Nothing- Right (value, "") -> Just value- Right (_value, _) -> Nothing-{-# INLINE readInt #-}--readByteStringInt :: (HasCallStack) => C.ByteString -> Maybe Int-readByteStringInt s = case C.readInt (C.strip s) of- Nothing -> Nothing- Just (value, "") -> Just value- Just (_value, _) -> Nothing-{-# INLINE readByteStringInt #-}--readByteStringDouble :: (HasCallStack) => C.ByteString -> Maybe Double-readByteStringDouble s =- let- readFunc = if C.any (\c -> c == 'e' || c == 'E') s then readExponential else readDecimal- in- case readSigned readFunc (C.strip s) of- Nothing -> Nothing- Just (value, "") -> Just value- Just (_value, _) -> Nothing-{-# INLINE readByteStringDouble #-}--readDouble :: (HasCallStack) => T.Text -> Maybe Double-readDouble s =- case signed double s of- Left _ -> Nothing- Right (value, "") -> Just value- Right (_value, _) -> Nothing-{-# INLINE readDouble #-}--readIntegerEither :: (HasCallStack) => T.Text -> Either T.Text Integer-readIntegerEither s = case signed decimal (T.strip s) of- Left _ -> Left s- Right (value, "") -> Right value- Right (_value, _) -> Left s-{-# INLINE readIntegerEither #-}--readIntEither :: (HasCallStack) => T.Text -> Either T.Text Int-readIntEither s = case signed decimal (T.strip s) of- Left _ -> Left s- Right (value, "") -> Right value- Right (_value, _) -> Left s-{-# INLINE readIntEither #-}--readDoubleEither :: (HasCallStack) => T.Text -> Either T.Text Double-readDoubleEither s =- case signed double s of- Left _ -> Left s- Right (value, "") -> Right value- Right (_value, _) -> Left s-{-# INLINE readDoubleEither #-}---- ------------------------------------------------------------------------------ Attoparsec CSV parser combinators (shared between Lazy.IO.CSV and others)--- -----------------------------------------------------------------------------parseSep :: Char -> T.Text -> [T.Text]-parseSep c s = either error id (parseOnly (record c) s)-{-# INLINE parseSep #-}--record :: Char -> Parser [T.Text]-record c =- field c `sepBy1` char c- <?> "record"-{-# INLINE record #-}--parseRow :: Char -> Parser [T.Text]-parseRow c = (record c <* lineEnd) <?> "record-new-line"--field :: Char -> Parser T.Text-field c =- quotedField <|> unquotedField c- <?> "field"-{-# INLINE field #-}--unquotedTerminators :: Char -> S.Set Char-unquotedTerminators sep = S.fromList [sep, '\n', '\r', '"']--unquotedField :: Char -> Parser T.Text-unquotedField sep =- takeWhile (not . (`S.member` terminators)) <?> "unquoted field"- where- terminators = unquotedTerminators sep-{-# INLINE unquotedField #-}--quotedField :: Parser T.Text-quotedField = char '"' *> contents <* char '"' <?> "quoted field"- where- contents = fold <$> many (unquote <|> unescape)- where- unquote = takeWhile1 (notInClass "\"\\")- unescape =- char '\\' *> do- T.singleton <$> do- char '\\' <|> char '"'-{-# INLINE quotedField #-}--lineEnd :: Parser ()-lineEnd =- (endOfLine <|> endOfInput)- <?> "end of line"-{-# INLINE lineEnd #-}---- | First pass to count rows for exact allocation.-countRows :: Char -> FilePath -> IO Int-countRows c path = withFile path ReadMode $! go 0 ""- where- go n input h = do- isEOF <- hIsEOF h- if isEOF && input == mempty- then pure n- else- parseWith (TIO.hGetChunk h) (parseRow c) input >>= \case- Fail unconsumed ctx er -> do- erpos <- hTell h- fail $- "Failed to parse CSV file around "- <> show erpos- <> " byte; due: "- <> show er- <> "; context: "- <> show ctx- <> " "- <> show unconsumed- Partial _ -> fail $ "Partial handler is called; n = " <> show n- Done (unconsumed :: T.Text) _ ->- go (n + 1) unconsumed h-{-# INLINE countRows #-}---- | Infer the Haskell type name from a text sample.-inferValueType :: T.Text -> T.Text-inferValueType s = case readInt s of- Just _ -> "Int"- Nothing -> case readDouble s of- Just _ -> "Double"- Nothing -> "Other"-{-# INLINE inferValueType #-}---- | Read a single CSV row from a handle using the given separator.-readSingleLine :: Char -> T.Text -> Handle -> IO ([T.Text], T.Text)-readSingleLine c unused handle =- parseWith (TIO.hGetChunk handle) (parseRow c) unused >>= \case- Fail _unconsumed ctx er -> do- erpos <- hTell handle- fail $- "Failed to parse CSV file around "- <> show erpos- <> " byte; due: "- <> show er- <> "; context: "- <> show ctx- Partial _ -> fail "Partial handler is called"- Done (unconsumed :: T.Text) (row :: [T.Text]) ->- return (row, unconsumed)
− src/DataFrame/Internal/Row.hs
@@ -1,175 +0,0 @@-{-# LANGUAGE ExistentialQuantification #-}-{-# LANGUAGE ExplicitNamespaces #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE InstanceSigs #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}--module DataFrame.Internal.Row where--import qualified Data.List as L-import qualified Data.Map as M-import qualified Data.Text as T-import qualified Data.Vector as V-import qualified Data.Vector.Unboxed as VU--import Control.Exception (throw)-import Data.Function (on)-import Data.Maybe (fromMaybe)-import Data.Type.Equality (TestEquality (..))-import Data.Typeable (type (:~:) (..))-import DataFrame.Errors (DataFrameException (..))-import DataFrame.Internal.Column-import DataFrame.Internal.DataFrame-import DataFrame.Internal.Expression (Expr (..))-import Type.Reflection (typeOf, typeRep)--data Any where- Value :: (Columnable a) => a -> Any--instance Eq Any where- (==) :: Any -> Any -> Bool- (Value a) == (Value b) = fromMaybe False $ do- Refl <- testEquality (typeOf a) (typeOf b)- return $ a == b--instance Show Any where- show :: Any -> String- show (Value a) = T.unpack (showValue a)--showValue :: forall a. (Columnable a) => a -> T.Text-showValue v = case testEquality (typeRep @a) (typeRep @T.Text) of- Just Refl -> v- Nothing -> case testEquality (typeRep @a) (typeRep @String) of- Just Refl -> T.pack v- Nothing -> (T.pack . show) v---- | Wraps a value into an \Any\ type. This helps up represent rows as heterogenous lists.-toAny :: forall a. (Columnable a) => a -> Any-toAny = Value---- | Unwraps a value from an \Any\ type.-fromAny :: forall a. (Columnable a) => Any -> Maybe a-fromAny (Value (v :: b)) = do- Refl <- testEquality (typeRep @a) (typeRep @b)- pure v--type Row = V.Vector Any--(!?) :: [a] -> Int -> Maybe a-(!?) [] _ = Nothing-(!?) (x : _) 0 = Just x-(!?) (_x : xs) n = (!?) xs (n - 1)--mkColumnFromRow :: Int -> [[Any]] -> Column-mkColumnFromRow i rows = case rows of- [] -> fromList ([] :: [T.Text])- (row : _) -> case row !? i of- Nothing -> fromList ([] :: [T.Text])- Just (Value (v :: a)) -> fromList $ reverse $ L.foldl' addToList [v] (drop 1 rows)- where- addToList acc r = case r !? i of- Nothing -> acc- Just (Value (v' :: b)) -> case testEquality (typeRep @a) (typeRep @b) of- Nothing -> acc- Just Refl -> v' : acc--{- | Converts the entire dataframe to a list of rows.--Each row contains all columns in the dataframe, ordered by their column indices.-The rows are returned in their natural order (from index 0 to n-1).--==== __Examples__-->>> toRowList df-[[("name", "Alice"), ("age", 25), ...], [("name", "Bob"), ("age", 30), ...], ...]--==== __Performance note__--This function materializes all rows into a list, which may be memory-intensive-for large dataframes. Consider using 'toRowVector' if you need random access-or streaming operations.--}-toRowList :: DataFrame -> [[(T.Text, Any)]]-toRowList df =- let- names = map fst (L.sortBy (compare `on` snd) $ M.toList (columnIndices df))- in- map- (zip names . V.toList . mkRowRep df names)- [0 .. (fst (dataframeDimensions df) - 1)]--{- | Converts the dataframe to a vector of rows with only the specified columns.--Each row will contain only the columns named in the @names@ parameter.-This is useful when you only need a subset of columns or want to control-the column order in the resulting rows.--==== __Parameters__--[@names@] List of column names to include in each row. The order of names- determines the order of fields in the resulting rows.--[@df@] The dataframe to convert.--==== __Examples__-->>> toRowVector ["name", "age"] df-Vector of rows with only name and age fields-->>> toRowVector [] df -- Empty column list-Vector of empty rows (one per dataframe row)--}-toRowVector :: [T.Text] -> DataFrame -> V.Vector Row-toRowVector names df = V.generate (fst (dataframeDimensions df)) (mkRowRep df names)--{- | Given a row gets the value associated with a field.--==== __Examples__-->>> map (rowValue (F.col @Int "age")) (toRowList df)-[25,30, ...]--}-rowValue :: forall a. Expr a -> [(T.Text, Any)] -> Maybe a-rowValue (Col name) row = lookup name row >>= fromAny @a-rowValue _ _ = error "Can only get rowValue of column reference"--mkRowFromArgs :: [T.Text] -> DataFrame -> Int -> Row-mkRowFromArgs names df i = V.map get (V.fromList names)- where- get name = case getColumn name df of- Nothing ->- throw $- ColumnsNotFoundException- [name]- "[INTERNAL] mkRowFromArgs"- (M.keys $ columnIndices df)- Just (BoxedColumn _ column) -> toAny (column V.! i)- Just (UnboxedColumn _ column) -> toAny (column VU.! i)---- This function will return the items in the order that is specified--- by the user. For example, if the dataframe consists of the columns--- "Age", "Pclass", "Name", and the user asks for ["Name", "Age"],--- this will order the values in the order ["Mr Smith", 50]-mkRowRep :: DataFrame -> [T.Text] -> Int -> Row-mkRowRep df names i = V.generate (L.length names) (\index -> get (names' V.! index))- where- names' = V.fromList names- throwError name =- error $- "Column "- ++ T.unpack name- ++ " has less items than "- ++ "the other columns at index "- ++ show i- get name = case getColumn name df of- Just (BoxedColumn _ c) -> case c V.!? i of- Just e -> toAny e- Nothing -> throwError name- Just (UnboxedColumn _ c) -> case c VU.!? i of- Just e -> toAny e- Nothing -> throwError name- Nothing ->- throw $ ColumnsNotFoundException [name] "mkRowRep" (M.keys $ columnIndices df)
− src/DataFrame/Internal/Schema.hs
@@ -1,242 +0,0 @@-{-# LANGUAGE AllowAmbiguousTypes #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE GADTs #-}-{-# 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--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.-data SchemaType where- -- | Constructor carrying a 'Proxy' of the element type.- SType :: (Columnable a, Read a) => P.Proxy a -> SchemaType--{- | Show the underlying element type using 'typeRep'.--==== __Examples__->>> :set -XTypeApplications->>> show (schemaType @Bool)-"Bool"--}-instance Show SchemaType where- show :: SchemaType -> String- show (SType (_ :: P.Proxy a)) = show (typeRep @a)--{- | Two 'SchemaType's are equal iff their element types are the same.--==== __Examples__->>> :set -XTypeApplications->>> schemaType @Int == schemaType @Int-True-->>> schemaType @Int == schemaType @Integer-False--}-instance Eq SchemaType where- (==) :: SchemaType -> SchemaType -> Bool- (==) (SType (_ :: P.Proxy a)) (SType (_ :: P.Proxy b)) =- isJust (testEquality (typeRep @a) (typeRep @b))--{- | Construct a 'SchemaType' for the given @a@.--==== __Examples__->>> :set -XTypeApplications->>> schemaType @T.Text == schemaType @T.Text-True-->>> show (schemaType @Double)-"Double"--}-schemaType :: forall a. (Columnable a, Read a) => SchemaType-schemaType = SType (P.Proxy @a)--{- | Logical schema of a 'DataFrame': a mapping from column names to their-element types ('SchemaType').--==== __Examples__-Constructing and querying a schema:-->>> import qualified Data.Map as M->>> import qualified Data.Text as T->>> let s = Schema (M.fromList [("country", schemaType @T.Text), ("amount", schemaType @Double)])->>> M.lookup "amount" (elements s) == Just (schemaType @Double)-True--Extending a schema:-->>> let s' = Schema (M.insert "discount" (schemaType @Double) (elements s))->>> M.member "discount" (elements s')-True--Equality is structural over the map contents:-->>> let a = Schema (M.fromList [("x", schemaType @Int), ("y", schemaType @Double)])->>> let b = Schema (M.fromList [("y", schemaType @Double), ("x", schemaType @Int)])->>> a == b-True--}-newtype Schema = Schema- { elements :: M.Map T.Text SchemaType- {- ^ Mapping from /column name/ to its 'SchemaType'.-- Invariant: keys are unique column names. A missing key means the column- is not present in the schema.- -}- }- deriving (Show, Eq)--{- | Construct a 'Schema' from a list of @(columnName, schemaType)@ pairs.--==== __Example__->>> :set -XTypeApplications->>> import qualified Data.Text as T->>> let s = makeSchema [("name", schemaType @T.Text), ("age", schemaType @Int)]->>> M.member "age" (elements s)-True--}-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/Internal/Statistics.hs
@@ -1,285 +0,0 @@-{-# LANGUAGE BangPatterns #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE ScopedTypeVariables #-}--module DataFrame.Internal.Statistics where--import qualified Data.Vector as V-import qualified Data.Vector.Algorithms.Intro as VA-import qualified Data.Vector.Mutable as VM-import qualified Data.Vector.Unboxed as VU-import qualified Data.Vector.Unboxed.Mutable as VUM--import Control.Exception (throw)-import Control.Monad.ST (runST)-import DataFrame.Errors (DataFrameException (..))--mean' :: (Real a, VU.Unbox a) => VU.Vector a -> Double-mean' samp- | VU.null samp = throw $ EmptyDataSetException "mean"- | otherwise = rtf (VU.sum samp) / fromIntegral (VU.length samp)-{-# INLINE [0] mean' #-}--meanDouble' :: VU.Vector Double -> Double-meanDouble' samp- | VU.null samp = throw $ EmptyDataSetException "mean"- | otherwise = VU.sum samp / fromIntegral (VU.length samp)-{-# INLINE meanDouble' #-}--meanInt' :: VU.Vector Int -> Double-meanInt' samp- | VU.null samp = throw $ EmptyDataSetException "mean"- | otherwise = fromIntegral (VU.sum samp) / fromIntegral (VU.length samp)-{-# INLINE meanInt' #-}--{-# RULES-"mean'/Double" [1] forall (xs :: VU.Vector Double).- mean' xs =- meanDouble' xs-"mean'/Int" [1] forall (xs :: VU.Vector Int).- mean' xs =- meanInt' xs- #-}--median' :: (Real a, VU.Unbox a) => VU.Vector a -> Double-median' samp- | VU.null samp = throw $ EmptyDataSetException "median"- | otherwise = runST $ do- mutableSamp <- VU.thaw samp- VA.sort mutableSamp- let len = VU.length samp- middleIndex = len `div` 2- middleElement <- VUM.read mutableSamp middleIndex- if odd len- then pure (rtf middleElement)- else do- prev <- VUM.read mutableSamp (middleIndex - 1)- pure (rtf (middleElement + prev) / 2)-{-# INLINE median' #-}---- accumulator: count, mean, m2-data VarAcc- = VarAcc {-# UNPACK #-} !Int {-# UNPACK #-} !Double {-# UNPACK #-} !Double- deriving (Show)--varianceStep :: VarAcc -> Double -> VarAcc-varianceStep (VarAcc !n !meanVal !m2) !x =- let !n' = n + 1- !delta = x - meanVal- !meanVal' = meanVal + delta / fromIntegral n'- !m2' = m2 + delta * (x - meanVal')- in VarAcc n' meanVal' m2'-{-# INLINE varianceStep #-}--computeVariance :: VarAcc -> Double-computeVariance (VarAcc !n _ !m2)- | n < 2 = 0 -- or error "variance of <2 samples"- | otherwise = m2 / fromIntegral (n - 1)-{-# INLINE computeVariance #-}--variance' :: (Real a, VU.Unbox a) => VU.Vector a -> Double-variance' = computeVariance . VU.foldl' varianceStep (VarAcc 0 0 0) . VU.map rtf-{-# INLINE variance' #-}--varianceDouble' :: VU.Vector Double -> Double-varianceDouble' = computeVariance . VU.foldl' varianceStep (VarAcc 0 0 0)-{-# INLINE varianceDouble' #-}---- accumulator: count, mean, m2, m3-data SkewAcc = SkewAcc !Int !Double !Double !Double deriving (Show)--skewnessStep :: (VU.Unbox a, Num a, Real a) => SkewAcc -> a -> SkewAcc-skewnessStep (SkewAcc !n !meanVal !m2 !m3) !x' =- let !n' = n + 1- x = rtf x'- !k = fromIntegral n'- !delta = x - meanVal- !meanVal' = meanVal + delta / k- !m2' = m2 + (delta ^ (2 :: Int) * (k - 1)) / k- !m3' =- m3- + (delta ^ (3 :: Int) * (k - 1) * (k - 2)) / k ^ (2 :: Int)- - (3 * delta * m2) / k- in SkewAcc n' meanVal' m2' m3'-{-# INLINE skewnessStep #-}--computeSkewness :: SkewAcc -> Double-computeSkewness (SkewAcc n _ m2 m3)- | n < 3 = 0 -- or error "skewness of <3 samples"- | otherwise = (sqrt (fromIntegral n - 1) * m3) / sqrt (m2 ^ (3 :: Int))-{-# INLINE computeSkewness #-}--skewness' :: (VU.Unbox a, Real a, Num a) => VU.Vector a -> Double-skewness' = computeSkewness . VU.foldl' skewnessStep (SkewAcc 0 0 0 0)-{-# INLINE skewness' #-}--data CorrelationStats- = CorrelationStats- {-# UNPACK #-} !Double- {-# UNPACK #-} !Double- {-# UNPACK #-} !Double- {-# UNPACK #-} !Double- {-# UNPACK #-} !Double--correlation' :: VU.Vector Double -> VU.Vector Double -> Maybe Double-correlation' xs ys- | n < 2 = Nothing- | VU.length xs /= VU.length ys = Nothing- | otherwise =- let nf = fromIntegral n- initial = CorrelationStats 0 0 0 0 0- (CorrelationStats sumX sumY sumXX sumYY sumXY) = VU.ifoldl' step initial xs-- !num = nf * sumXY - sumX * sumY- !den = sqrt ((nf * sumXX - sumX * sumX) * (nf * sumYY - sumY * sumY))- in Just (num / den)- where- n = VU.length xs- step (CorrelationStats sx sy sxx syy sxy) i x =- let !y = VU.unsafeIndex ys i- in CorrelationStats (sx + x) (sy + y) (sxx + x * x) (syy + y * y) (sxy + x * y)-{-# INLINE correlation' #-}--quantiles' ::- (VU.Unbox a, Num a, Real a) =>- VU.Vector Int -> Int -> VU.Vector a -> VU.Vector Double-quantiles' qs q samp- | VU.null samp = throw $ EmptyDataSetException "quantiles"- | q < 2 = throw $ WrongQuantileNumberException q- | VU.any (\i -> i < 0 || i > q) qs = throw $ WrongQuantileIndexException qs q- | otherwise = runST $ do- let !n = VU.length samp- mutableSamp <- VU.thaw samp- VA.sort mutableSamp- VU.mapM- ( \i -> do- let !p = fromIntegral i / fromIntegral q- !position = p * fromIntegral (n - 1) :: Double- !index = floor position :: Int- !f = position - fromIntegral index- x <- fmap rtf (VUM.read mutableSamp index)- if f == 0- then return x- else do- y <- fmap rtf (VUM.read mutableSamp (index + 1))- return $ (1 - f) * x + f * y- )- qs-{-# INLINE quantiles' #-}--percentile' :: (VU.Unbox a, Num a, Real a) => Int -> VU.Vector a -> Double-percentile' n = VU.head . quantiles' (VU.fromList [n]) 100--quantilesOrd' ::- (Ord a, Eq a) =>- VU.Vector Int -> Int -> V.Vector a -> V.Vector a-quantilesOrd' qs q samp- | V.null samp = throw $ EmptyDataSetException "quantiles"- | q < 2 = throw $ WrongQuantileNumberException q- | VU.any (\i -> i < 0 || i > q) qs = throw $ WrongQuantileIndexException qs q- | otherwise = runST $ do- let !n = V.length samp- mutableSamp <- V.thaw samp- VA.sort mutableSamp- V.mapM- ( \i -> do- let !p = fromIntegral i / fromIntegral q :: Double- !position = p * fromIntegral (n - 1)- !index = floor position :: Int- -- This is not exact for Ord instances.- -- Figure out how to make it so.- VM.read mutableSamp index- )- (V.convert qs)--percentileOrd' :: (Ord a, Eq a) => Int -> V.Vector a -> a-percentileOrd' n = V.head . quantilesOrd' (VU.fromList [n]) 100--interQuartileRange' :: (VU.Unbox a, Num a, Real a) => VU.Vector a -> Double-interQuartileRange' samp =- let quartiles = quantiles' (VU.fromList [1, 3]) 4 samp- in quartiles VU.! 1 - quartiles VU.! 0-{-# INLINE interQuartileRange' #-}--meanSquaredError :: VU.Vector Double -> VU.Vector Double -> Maybe Double-meanSquaredError target prediction =- let- squareDiff = VU.ifoldl' (\sq i e -> (e - target VU.! i) ^ (2 :: Int) + sq) 0 prediction- in- Just $ squareDiff / fromIntegral (max (VU.length target) (VU.length prediction))-{-# INLINE meanSquaredError #-}--mutualInformationBinned ::- Int -> VU.Vector Double -> VU.Vector Double -> Maybe Double-mutualInformationBinned k xs ys- | VU.length xs /= VU.length ys = Nothing- | VU.null xs = Nothing- | k < 2 = Nothing- | rx <= 0 || ry <= 0 = Just 0- | otherwise =- let bx = VU.map (binIndex xmin xmax k) xs- by = VU.map (binIndex ymin ymax k) ys- n = fromIntegral (VU.length xs) :: Double- mx = bincount k bx- my = bincount k by- mxy = jointBincount k bx by- in Just $- sum- [ let !cxy = fromIntegral c- !pxy = cxy / n- !px = fromIntegral (mx VU.! i) / n- !py = fromIntegral (my VU.! j) / n- in if c == 0 then 0 else pxy * logBase 2 (pxy / (px * py))- | i <- [0 .. k - 1]- , j <- [0 .. k - 1]- , let !c = mxy VU.! (i * k + j)- ]- where- (xmin, xmax) = (VU.minimum xs, VU.maximum xs)- (ymin, ymax) = (VU.minimum ys, VU.maximum ys)- rx = xmax - xmin- ry = ymax - ymin--binIndex :: Double -> Double -> Int -> Double -> Int-binIndex lo hi k x- | hi == lo = 0- | otherwise =- let !t = (x - lo) / (hi - lo)- !ix = floor (fromIntegral k * t) :: Int- in max 0 (min (k - 1) ix)-{-# INLINE binIndex #-}--bincount :: Int -> VU.Vector Int -> VU.Vector Int-bincount k bs = VU.create $ do- mv <- VU.thaw (VU.replicate k 0)- VU.forM_ bs $ \b -> do- let i- | b < 0 = 0- | b >= k = k - 1- | otherwise = b- x <- VUM.read mv i- VUM.write mv i (x + 1)- pure mv-{-# INLINE bincount #-}--jointBincount :: Int -> VU.Vector Int -> VU.Vector Int -> VU.Vector Int-jointBincount k bx by = VU.create $ do- mv <- VU.thaw (VU.replicate (k * k) 0)- VU.forM_ (VU.zip bx by) $ \(i, j) -> do- let ii = clamp i 0 (k - 1)- jj = clamp j 0 (k - 1)- ix = ii * k + jj- x <- VUM.read mv ix- VUM.write mv ix (x + 1)- pure mv- where- clamp z a b = max a (min b z)-{-# INLINE jointBincount #-}--rtf :: (Real a) => a -> Double-rtf = realToFrac-{-# NOINLINE [1] rtf #-}--{-# RULES-"rtf/Double" [2] forall (x :: Double). rtf x = x- #-}
− src/DataFrame/Internal/Types.hs
@@ -1,153 +0,0 @@-{-# LANGUAGE AllowAmbiguousTypes #-}-{-# LANGUAGE ConstraintKinds #-}-{-# LANGUAGE DataKinds #-}-{-# LANGUAGE ExistentialQuantification #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE PolyKinds #-}-{-# LANGUAGE RankNTypes #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}-{-# LANGUAGE TypeFamilies #-}-{-# LANGUAGE UndecidableInstances #-}--module DataFrame.Internal.Types where--import Data.Int (Int16, Int32, Int64, Int8)-import Data.Kind (Constraint, Type)-import Data.Typeable (Typeable)-import qualified Data.Vector.Unboxed as VU-import Data.Word (Word16, Word32, Word64, Word8)--type Columnable' a = (Typeable a, Show a, Eq a)--{- | A type with column representations used to select the-"right" representation when specializing the `toColumn` function.--}-data Rep- = RBoxed- | RUnboxed- | RNullableBoxed---- | Type-level if statement.-type family If (cond :: Bool) (yes :: k) (no :: k) :: k where- If 'True yes _ = yes- If 'False _ no = no---- | All unboxable types (according to the `vector` package).-type family Unboxable (a :: Type) :: Bool where- Unboxable Int = 'True- Unboxable Int8 = 'True- Unboxable Int16 = 'True- Unboxable Int32 = 'True- Unboxable Int64 = 'True- Unboxable Word = 'True- Unboxable Word8 = 'True- Unboxable Word16 = 'True- Unboxable Word32 = 'True- Unboxable Word64 = 'True- Unboxable Char = 'True- Unboxable Bool = 'True- Unboxable Double = 'True- Unboxable Float = 'True- Unboxable _ = 'False--type family Numeric (a :: Type) :: Bool where- Numeric Integer = 'True- Numeric Int = 'True- Numeric Int8 = 'True- Numeric Int16 = 'True- Numeric Int32 = 'True- Numeric Int64 = 'True- Numeric Word = 'True- Numeric Word8 = 'True- Numeric Word16 = 'True- Numeric Word32 = 'True- Numeric Word64 = 'True- Numeric Double = 'True- Numeric Float = 'True- Numeric _ = 'False---- | Compute the column representation tag for any 'a'.-type family KindOf a :: Rep where- KindOf (Maybe a) = 'RNullableBoxed- KindOf a = If (Unboxable a) 'RUnboxed 'RBoxed---- | Type-level boolean for constraint/type comparison.-data SBool (b :: Bool) where- STrue :: SBool 'True- SFalse :: SBool 'False---- | The runtime witness for our type-level branching.-class SBoolI (b :: Bool) where- sbool :: SBool b--instance SBoolI 'True where sbool = STrue-instance SBoolI 'False where sbool = SFalse---- | Type-level function to determine whether or not a type is unboxa-sUnbox :: forall a. (SBoolI (Unboxable a)) => SBool (Unboxable a)-sUnbox = sbool @(Unboxable a)--sNumeric :: forall a. (SBoolI (Numeric a)) => SBool (Numeric a)-sNumeric = sbool @(Numeric a)--type family When (flag :: Bool) (c :: Constraint) :: Constraint where- When 'True c = c- When 'False c = () -- empty constraint--type UnboxIf a = When (Unboxable a) (VU.Unbox a)--type family IntegralTypes (a :: Type) :: Bool where- IntegralTypes Integer = 'True- IntegralTypes Int = 'True- IntegralTypes Int8 = 'True- IntegralTypes Int16 = 'True- IntegralTypes Int32 = 'True- IntegralTypes Int64 = 'True- IntegralTypes Word = 'True- IntegralTypes Word8 = 'True- IntegralTypes Word16 = 'True- IntegralTypes Word32 = 'True- IntegralTypes Word64 = 'True- IntegralTypes _ = 'False--sIntegral :: forall a. (SBoolI (IntegralTypes a)) => SBool (IntegralTypes a)-sIntegral = sbool @(IntegralTypes a)--type IntegralIf a = When (IntegralTypes a) (Integral a)--type family FloatingTypes (a :: Type) :: Bool where- FloatingTypes Float = 'True- FloatingTypes Double = 'True- FloatingTypes _ = 'False--sFloating :: forall a. (SBoolI (FloatingTypes a)) => SBool (FloatingTypes a)-sFloating = sbool @(FloatingTypes a)--type FloatingIf a = When (FloatingTypes a) (Real a, Fractional a)--{- | Numeric type promotion: resolves the common type for mixed arithmetic.-Double dominates over Float/Int; Float dominates over Int; same types stay unchanged.--}-type family Promote (a :: Type) (b :: Type) :: Type where- Promote a a = a- Promote Double _ = Double- Promote _ Double = Double- Promote Float _ = Float- Promote _ Float = Float- Promote Int64 _ = Int64- Promote _ Int64 = Int64- Promote Int32 _ = Int32- Promote _ Int32 = Int32- Promote a _ = a--{- | Like 'Promote', but integral × integral → Double for use with './' .-Double\/Float still dominate; any two integral types (same or mixed) become Double.--}-type family PromoteDiv (a :: Type) (b :: Type) :: Type where- PromoteDiv Double _ = Double- PromoteDiv _ Double = Double- PromoteDiv Float _ = Float- PromoteDiv _ Float = Float- PromoteDiv _ _ = Double -- Int/Int32/Int64 in any combination
− src/DataFrame/Lazy.hs
@@ -1,3 +0,0 @@-module DataFrame.Lazy (module DataFrame.Lazy.Internal.DataFrame) where--import DataFrame.Lazy.Internal.DataFrame
− src/DataFrame/Lazy/IO/Binary.hs
@@ -1,413 +0,0 @@-{-# LANGUAGE ExplicitNamespaces #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE StrictData #-}-{-# LANGUAGE TypeApplications #-}--{- | Simple column-oriented binary spill format (DFBN).--Layout (all integers little-endian):--@-[magic: 4 bytes] "DFBN"-[num_columns: 4 bytes] Word32- per column:- [name_len: 2 bytes] Word16 (byte length of UTF-8 name)- [name: name_len bytes]- [type_tag: 1 byte] Word8-[num_rows: 8 bytes] Word64--per column data block (order matches schema):- type_tag 0 (Int): num_rows × Int64 LE- type_tag 1 (Double): num_rows × Double LE (IEEE 754)- type_tag 2 (Text): (num_rows+1) × Word32 offsets ++ payload bytes (UTF-8)- type_tag 3 (Maybe Int): ceil(num_rows/8)-byte null bitmap ++ num_rows × Int64 LE- type_tag 4 (Maybe Double): ceil(num_rows/8)-byte null bitmap ++ num_rows × Double LE- type_tag 5 (Maybe Text): ceil(num_rows/8)-byte null bitmap- ++ (num_rows+1) × Word32 offsets ++ payload bytes-@--Null bitmap: bit @i@ of byte @i\/8@ is 1 when row @i@ is non-null.--}-module DataFrame.Lazy.IO.Binary (- spillToDisk,- readSpilled,- withSpilled,-) where--import Control.Exception (SomeException, bracket, try)-import Control.Monad (foldM, void, when)-import qualified Data.ByteString as BS-import qualified Data.ByteString.Builder as BSB-import qualified Data.ByteString.Internal as BSI-import qualified Data.ByteString.Unsafe as BSU-import qualified Data.List as L-import qualified Data.Map.Strict as M-import qualified Data.Text as T-import qualified Data.Text.Encoding as TE-import qualified Data.Vector as V-import qualified Data.Vector.Storable as VS-import qualified Data.Vector.Unboxed as VU--import Data.Bits (setBit, shiftL, testBit, (.|.))-import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl))-import Data.Word (Word16, Word32, Word64, Word8)-import qualified DataFrame.Internal.Binary as Binary-import DataFrame.Internal.Column (- Column (..),- bitmapTestBit,- buildBitmapFromValid,- )-import DataFrame.Internal.DataFrame (DataFrame (..))-import Foreign (ForeignPtr, castForeignPtr, plusForeignPtr, sizeOf)-import System.Directory (getTemporaryDirectory, removeFile)-import System.IO (IOMode (..), hClose, openTempFile, withFile)-import Type.Reflection (typeRep)---- ------------------------------------------------------------------------------ Type tags--- -----------------------------------------------------------------------------tagInt, tagDouble, tagText, tagMaybeInt, tagMaybeDouble, tagMaybeText :: Word8-tagInt = 0-tagDouble = 1-tagText = 2-tagMaybeInt = 3-tagMaybeDouble = 4-tagMaybeText = 5---- ------------------------------------------------------------------------------ Write--- ------------------------------------------------------------------------------- | Serialise a 'DataFrame' to a DFBN binary file.-spillToDisk :: FilePath -> DataFrame -> IO ()-spillToDisk path df =- withFile path WriteMode $ \h -> BSB.hPutBuilder h (buildDataFrame df)--buildDataFrame :: DataFrame -> BSB.Builder-buildDataFrame df =- BSB.byteString "DFBN"- <> BSB.word32LE ncols- <> foldMap (uncurry buildColumnSchema) (zip names cols)- <> BSB.word64LE nrows- <> foldMap (buildColumnData nrowsInt) cols- where- names =- fmap- fst- (L.sortBy (\a b -> compare (snd a) (snd b)) (M.toList (columnIndices df)))- ncols = fromIntegral (length names) :: Word32- cols = V.toList (columns df)- nrowsInt = fst (dataframeDimensions df)- nrows = fromIntegral nrowsInt :: Word64--buildColumnSchema :: T.Text -> Column -> BSB.Builder-buildColumnSchema name col =- BSB.word16LE nameLen- <> BSB.byteString nameBytes- <> BSB.word8 (columnTypeTag col)- where- nameBytes = TE.encodeUtf8 name- nameLen = fromIntegral (BS.length nameBytes) :: Word16--columnTypeTag :: Column -> Word8-columnTypeTag (UnboxedColumn Nothing (_ :: VU.Vector a)) =- case testEquality (typeRep @a) (typeRep @Int) of- Just Refl -> tagInt- Nothing -> case testEquality (typeRep @a) (typeRep @Double) of- Just Refl -> tagDouble- Nothing -> error "spillToDisk: unsupported UnboxedColumn element type"-columnTypeTag (UnboxedColumn (Just _) (_ :: VU.Vector a)) =- case testEquality (typeRep @a) (typeRep @Int) of- Just Refl -> tagMaybeInt- Nothing -> case testEquality (typeRep @a) (typeRep @Double) of- Just Refl -> tagMaybeDouble- Nothing -> error "spillToDisk: unsupported nullable UnboxedColumn element type"-columnTypeTag (BoxedColumn Nothing _) = tagText-columnTypeTag (BoxedColumn (Just _) _) = tagMaybeText--buildColumnData :: Int -> Column -> BSB.Builder-buildColumnData _ (UnboxedColumn Nothing (v :: VU.Vector a)) =- case testEquality (typeRep @a) (typeRep @Int) of- Just Refl -> buildIntVector v- Nothing ->- case testEquality (typeRep @a) (typeRep @Double) of- Just Refl -> buildDoubleVector v- Nothing -> error "spillToDisk: unsupported UnboxedColumn element type"-buildColumnData _ (UnboxedColumn (Just bm) (v :: VU.Vector a)) =- case testEquality (typeRep @a) (typeRep @Int) of- Just Refl ->- buildNullBitmap (V.generate (VU.length v) (bitmapTestBit bm))- <> buildIntVector v- Nothing ->- case testEquality (typeRep @a) (typeRep @Double) of- Just Refl ->- buildNullBitmap (V.generate (VU.length v) (bitmapTestBit bm))- <> buildDoubleVector v- Nothing -> error "spillToDisk: unsupported nullable UnboxedColumn element type"-buildColumnData _ (BoxedColumn Nothing (v :: V.Vector a)) =- case testEquality (typeRep @a) (typeRep @T.Text) of- Just Refl -> buildTextVector v- Nothing -> error "spillToDisk: unsupported BoxedColumn element type"-buildColumnData _ (BoxedColumn (Just bm) (v :: V.Vector a)) =- let isValidVec = V.generate (V.length v) (bitmapTestBit bm)- showText x = case testEquality (typeRep @a) (typeRep @T.Text) of- Just Refl -> x- Nothing -> T.pack (show x)- texts = V.imap (\i x -> if bitmapTestBit bm i then showText x else T.empty) v- in buildNullBitmap isValidVec <> buildTextVector texts--{- | Bulk-encode an Int vector as 8-byte LE values (native layout on LE platforms).-hPutBuilder flushes synchronously so the underlying ForeignPtr outlives the Builder.--}-buildIntVector :: VU.Vector Int -> BSB.Builder-buildIntVector v =- let sv = VU.convert v :: VS.Vector Int- (fp, n) = VS.unsafeToForeignPtr0 sv- bs = BSI.fromForeignPtr (castForeignPtr fp) 0 (n * sizeOf (0 :: Int))- in BSB.byteString bs---- | Bulk-encode a Double vector as 8-byte LE IEEE 754 values (native layout on LE platforms).-buildDoubleVector :: VU.Vector Double -> BSB.Builder-buildDoubleVector v =- let sv = VU.convert v :: VS.Vector Double- (fp, n) = VS.unsafeToForeignPtr0 sv- bs = BSI.fromForeignPtr (castForeignPtr fp) 0 (n * sizeOf (0 :: Double))- in BSB.byteString bs---- | Write a Text vector: (num_rows+1) Word32 offsets followed by UTF-8 payload.-buildTextVector :: V.Vector T.Text -> BSB.Builder-buildTextVector v =- foldMap BSB.word32LE offsets <> foldMap BSB.byteString encoded- where- encoded = V.toList (V.map TE.encodeUtf8 v)- offsets = scanl (\acc bs -> acc + fromIntegral (BS.length bs)) (0 :: Word32) encoded---- | Build a null-validity bitmap: 1 bit per row, packed LSB-first into bytes.-buildNullBitmap :: V.Vector Bool -> BSB.Builder-buildNullBitmap valids = foldMap (BSB.word8 . mkByte) [0 .. numBytes - 1]- where- n = V.length valids- numBytes = (n + 7) `div` 8- mkByte byteIdx =- foldr- ( \bit acc ->- let row = byteIdx * 8 + bit- in if row < n && (valids V.! row) then setBit acc bit else acc- )- (0 :: Word8)- [0 .. 7]---- ------------------------------------------------------------------------------ Read--- ------------------------------------------------------------------------------- | @(new_offset, value)@-type ParseResult a = Either String (Int, a)---- | Deserialise a DFBN binary file into a 'DataFrame'.-readSpilled :: FilePath -> IO DataFrame-readSpilled path = do- bs <- BS.readFile path- case parseDataFrame bs 0 of- Left err -> fail ("readSpilled: " <> err)- Right (_, df) -> return df--parseDataFrame :: BS.ByteString -> Int -> ParseResult DataFrame-parseDataFrame bs off0 = do- (off1, magic) <- readBytes bs off0 4- when (magic /= "DFBN") $ Left "bad magic bytes"- (off2, ncols) <- readWord32LE bs off1- let ncolsInt = fromIntegral ncols :: Int- (off3, schema) <- readN ncolsInt (readColumnSchema bs) off2- (off4, nrows64) <- readWord64LE bs off3- let nrows = fromIntegral nrows64 :: Int- (off5, cols) <-- foldM- ( \(o, acc) (_, tag) -> do- (o', col) <- readColumnData bs o nrows tag- return (o', acc ++ [col])- )- (off4, [])- schema- let names = fmap fst schema- return- ( off5- , DataFrame- { columns = V.fromList cols- , columnIndices = M.fromList (zip names [0 ..])- , dataframeDimensions = (nrows, ncolsInt)- , derivingExpressions = M.empty- }- )--readColumnSchema :: BS.ByteString -> Int -> ParseResult (T.Text, Word8)-readColumnSchema bs off = do- (off1, nameLen) <- readWord16LE bs off- let nameLenInt = fromIntegral nameLen :: Int- (off2, nameBytes) <- readBytes bs off1 nameLenInt- (off3, tag) <- readWord8 bs off2- return (off3, (TE.decodeUtf8 nameBytes, tag))--readColumnData :: BS.ByteString -> Int -> Int -> Word8 -> ParseResult Column-readColumnData bs off nrows tag- | tag == tagInt = do- (off', v) <- readIntColumn bs off nrows- return (off', UnboxedColumn Nothing v)- | tag == tagDouble = do- (off', v) <- readDoubleColumn bs off nrows- return (off', UnboxedColumn Nothing v)- | tag == tagText = do- (off', v) <- readTextColumn bs off nrows- return (off', BoxedColumn Nothing v)- | tag == tagMaybeInt = do- (off1, bitmap) <- readNullBitmap bs off nrows- (off2, v) <- readIntColumn bs off1 nrows- let bm = buildBitmapFromValid (VU.fromList (map (\b -> if b then 1 else 0) bitmap))- return (off2, UnboxedColumn (Just bm) v)- | tag == tagMaybeDouble = do- (off1, bitmap) <- readNullBitmap bs off nrows- (off2, v) <- readDoubleColumn bs off1 nrows- let bm = buildBitmapFromValid (VU.fromList (map (\b -> if b then 1 else 0) bitmap))- return (off2, UnboxedColumn (Just bm) v)- | tag == tagMaybeText = do- (off1, bitmap) <- readNullBitmap bs off nrows- (off2, v) <- readTextColumn bs off1 nrows- let bm = buildBitmapFromValid (VU.fromList (map (\b -> if b then 1 else 0) bitmap))- return (off2, BoxedColumn (Just bm) v)- | otherwise = Left ("unknown type tag " <> show tag)--{- | Zero-copy Int column read: reuses the ByteString buffer's ForeignPtr.-Safe as long as 'bs' stays live during the caller's use of the resulting vector.-Only correct on little-endian platforms (aarch64/x86_64).--}-readIntColumn :: BS.ByteString -> Int -> Int -> ParseResult (VU.Vector Int)-readIntColumn bs off nrows- | off + nrows * 8 > BS.length bs = Left "unexpected end of input"- | otherwise =- let (fp, bsOff, _) = BSI.toForeignPtr bs- fp' = castForeignPtr (plusForeignPtr fp (bsOff + off)) :: ForeignPtr Int- sv = VS.unsafeFromForeignPtr0 fp' nrows :: VS.Vector Int- in Right (off + nrows * 8, VU.convert sv)--{- | Zero-copy Double column read: reuses the ByteString buffer's ForeignPtr.-Safe as long as 'bs' stays live during the caller's use of the resulting vector.-Only correct on little-endian platforms (aarch64/x86_64).--}-readDoubleColumn ::- BS.ByteString -> Int -> Int -> ParseResult (VU.Vector Double)-readDoubleColumn bs off nrows- | off + nrows * 8 > BS.length bs = Left "unexpected end of input"- | otherwise =- let (fp, bsOff, _) = BSI.toForeignPtr bs- fp' = castForeignPtr (plusForeignPtr fp (bsOff + off)) :: ForeignPtr Double- sv = VS.unsafeFromForeignPtr0 fp' nrows :: VS.Vector Double- in Right (off + nrows * 8, VU.convert sv)--readTextColumn :: BS.ByteString -> Int -> Int -> ParseResult (V.Vector T.Text)-readTextColumn bs off nrows = do- offsets <- readWord32Array bs off (nrows + 1)- let payloadStart = off + (nrows + 1) * 4- totalPayload = fromIntegral (last offsets) :: Int- when (payloadStart + totalPayload > BS.length bs) $- Left "unexpected end of input"- let sizes =- zipWith- (\a b -> fromIntegral b - fromIntegral a :: Int)- offsets- (drop 1 offsets)- texts =- zipWith- ( \o sz ->- TE.decodeUtf8- (BS.take sz (BS.drop (payloadStart + fromIntegral o) bs))- )- offsets- sizes- return (payloadStart + totalPayload, V.fromList texts)---- | Read @nrows@ null-bitmap bits (ceil(nrows\/8) bytes).-readNullBitmap :: BS.ByteString -> Int -> Int -> ParseResult [Bool]-readNullBitmap bs off nrows- | off + numBytes > BS.length bs = Left "unexpected end of input"- | otherwise =- Right- ( off + numBytes- , take- nrows- [ testBit (BSU.unsafeIndex bs (off + row `div` 8)) (row `mod` 8)- | row <- [0 ..]- ]- )- where- numBytes = (nrows + 7) `div` 8--readWord8 :: BS.ByteString -> Int -> ParseResult Word8-readWord8 bs off- | off >= BS.length bs = Left "unexpected end of input"- | otherwise = Right (off + 1, BSU.unsafeIndex bs off)--readWord16LE :: BS.ByteString -> Int -> ParseResult Word16-readWord16LE bs off- | off + 2 > BS.length bs = Left "unexpected end of input"- | otherwise =- let b0 = fromIntegral (BSU.unsafeIndex bs off) :: Word16- b1 = fromIntegral (BSU.unsafeIndex bs (off + 1)) :: Word16- in Right (off + 2, b0 .|. (b1 `shiftL` 8))--readWord32LE :: BS.ByteString -> Int -> ParseResult Word32-readWord32LE bs off- | off + 4 > BS.length bs = Left "unexpected end of input"- | otherwise = Right (off + 4, Binary.littleEndianWord32 (BS.drop off bs))--readWord64LE :: BS.ByteString -> Int -> ParseResult Word64-readWord64LE bs off- | off + 8 > BS.length bs = Left "unexpected end of input"- | otherwise = Right (off + 8, Binary.littleEndianWord64 (BS.drop off bs))---- | Read @n@ consecutive Word32LE values starting at offset @off@.-readWord32Array :: BS.ByteString -> Int -> Int -> Either String [Word32]-readWord32Array bs off n- | off + n * 4 > BS.length bs = Left "unexpected end of input"- | otherwise =- Right- [ let i = off + k * 4- in Binary.littleEndianWord32 (BS.drop i bs)- | k <- [0 .. n - 1]- ]---- | Read @n@ bytes from @bs@ at @off@.-readBytes :: BS.ByteString -> Int -> Int -> ParseResult BS.ByteString-readBytes bs off n- | off + n > BS.length bs = Left "unexpected end of input"- | otherwise = Right (off + n, BS.take n (BS.drop off bs))---- | Apply @f@ @n@ times sequentially, threading the offset.-readN :: Int -> (Int -> ParseResult a) -> Int -> ParseResult [a]-readN 0 _ off = Right (off, [])-readN n f off = do- (off', x) <- f off- (off'', xs) <- readN (n - 1) f off'- return (off'', x : xs)---- ------------------------------------------------------------------------------ Bracket helper--- -----------------------------------------------------------------------------{- | Spill a DataFrame to a temporary file, run an action with the path,-then delete the file even if the action throws.--}-withSpilled :: DataFrame -> (FilePath -> IO a) -> IO a-withSpilled df action = do- tmpDir <- getTemporaryDirectory- bracket- ( do- (path, h) <- openTempFile tmpDir "dataframe_spill.dfbn"- hClose h- spillToDisk path df- return path- )- (\path -> void (try (removeFile path) :: IO (Either SomeException ())))- action
− src/DataFrame/Lazy/IO/CSV.hs
@@ -1,469 +0,0 @@-{-# LANGUAGE BangPatterns #-}-{-# LANGUAGE ExplicitNamespaces #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE LambdaCase #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE RankNTypes #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}--module DataFrame.Lazy.IO.CSV where--import qualified Data.ByteString as BS-import qualified Data.Map as M-import qualified Data.Proxy as P-import qualified Data.Text as T-import qualified Data.Text.Encoding as TextEncoding-import qualified Data.Text.IO as TIO-import qualified Data.Vector as V-import qualified Data.Vector.Mutable as VM-import qualified Data.Vector.Unboxed.Mutable as VUM--import Control.Monad (forM_, unless, when, zipWithM_)-import Data.Attoparsec.Text (IResult (..), parseWith)-import Data.Char (intToDigit)-import Data.IORef-import Data.Maybe (fromMaybe, isJust)-import Data.Type.Equality (TestEquality (testEquality))-import Data.Word (Word8)-import DataFrame.Internal.Column (- Column (..),- MutableColumn (..),- columnLength,- ensureOptional,- freezeColumn',- freezeColumnEither,- writeColumn,- )-import DataFrame.Internal.DataFrame (DataFrame (..))-import DataFrame.Internal.Parsing-import DataFrame.Internal.Schema (Schema, SchemaType (..), elements)-import DataFrame.Operations.Typing (SafeReadMode (..), effectiveSafeRead)-import System.IO-import Type.Reflection-import Prelude hiding (takeWhile)---- | Record for CSV read options.-data ReadOptions = ReadOptions- { hasHeader :: Bool- , inferTypes :: Bool- , safeRead :: SafeReadMode- {- ^ Default 'SafeReadMode' for columns without an entry in- 'safeReadOverrides'.- -}- , safeReadOverrides :: [(T.Text, SafeReadMode)]- -- ^ Per-column 'SafeReadMode' overrides; takes precedence over 'safeRead'.- , rowRange :: !(Maybe (Int, Int)) -- (start, length)- , seekPos :: !(Maybe Integer)- , totalRows :: !(Maybe Int)- , leftOver :: !T.Text- , rowsRead :: !Int- }--{- | By default we assume the file has a header and we infer types on read.-'safeRead' starts as 'NoSafeRead' — set it to 'MaybeRead' to wrap columns as-@Maybe a@, or 'EitherRead' to wrap as @Either Text a@ preserving the raw text-of any rows that fail to parse. Use 'safeReadOverrides' to pick a different-mode for specific columns.--}-defaultOptions :: ReadOptions-defaultOptions =- ReadOptions- { hasHeader = True- , inferTypes = True- , safeRead = NoSafeRead- , safeReadOverrides = []- , rowRange = Nothing- , seekPos = Nothing- , totalRows = Nothing- , leftOver = ""- , rowsRead = 0- }--{- | Reads a CSV file from the given path.-Note this file stores intermediate temporary files-while converting the CSV from a row to a columnar format.--}-readCsv :: FilePath -> IO DataFrame-readCsv path = fst <$> readSeparated ',' defaultOptions path--{- | Reads a tab separated file from the given path.-Note this file stores intermediate temporary files-while converting the CSV from a row to a columnar format.--}-readTsv :: FilePath -> IO DataFrame-readTsv path = fst <$> readSeparated '\t' defaultOptions path---- | Reads a character separated file into a dataframe using mutable vectors.-readSeparated ::- Char -> ReadOptions -> FilePath -> IO (DataFrame, (Integer, T.Text, Int))-readSeparated c opts path = do- totalRows' <- case totalRows opts of- Nothing ->- countRows c path >>= \total -> if hasHeader opts then return (total - 1) else return total- Just n -> if hasHeader opts then return (n - 1) else return n- let (_, len') = case rowRange opts of- Nothing -> (0, totalRows')- Just (start, len'') -> (start, min len'' (totalRows' - rowsRead opts))- withFile path ReadMode $ \handle -> do- firstRow <- fmap T.strip . parseSep c <$> TIO.hGetLine handle- let columnNames =- if hasHeader opts- then fmap (T.filter (/= '\"')) firstRow- else fmap (T.singleton . intToDigit) [0 .. (length firstRow - 1)]- -- If there was no header rewind the file cursor.- unless (hasHeader opts) $ hSeek handle AbsoluteSeek 0-- currPos <- hTell handle- when (isJust $ seekPos opts) $- hSeek handle AbsoluteSeek (fromMaybe currPos (seekPos opts))-- -- Initialize mutable vectors for each column- let numColumns = length columnNames- let numRows = len'- -- Use this row to infer the types of the rest of the column.- (dataRow, remainder) <- readSingleLine c (leftOver opts) handle-- -- This array will track the indices of all null values for each column.- nullIndices <- VM.unsafeNew numColumns- VM.set nullIndices []- mutableCols <- VM.unsafeNew numColumns- getInitialDataVectors numRows mutableCols dataRow-- -- Read rows into the mutable vectors- (unconsumed, r) <-- fillColumns numRows c mutableCols nullIndices remainder handle-- -- Freeze the mutable vectors into immutable ones- nulls' <- V.unsafeFreeze nullIndices- let !columnNamesV = V.fromList columnNames- cols <-- V.mapM- (freezeColumn columnNamesV mutableCols nulls' opts)- (V.generate numColumns id)- pos <- hTell handle-- return- ( DataFrame- { columns = cols- , columnIndices = M.fromList (zip columnNames [0 ..])- , dataframeDimensions = (maybe 0 columnLength (cols V.!? 0), V.length cols)- , derivingExpressions = M.empty- }- , (pos, unconsumed, r + 1)- )-{-# INLINE readSeparated #-}--getInitialDataVectors :: Int -> VM.IOVector MutableColumn -> [T.Text] -> IO ()-getInitialDataVectors n mCol xs = do- forM_ (zip [0 ..] xs) $ \(i, x) -> do- col <- case inferValueType x of- "Int" ->- MUnboxedColumn- <$> ( (VUM.unsafeNew n :: IO (VUM.IOVector Int)) >>= \c -> VUM.unsafeWrite c 0 (fromMaybe 0 $ readInt x) >> return c- )- "Double" ->- MUnboxedColumn- <$> ( (VUM.unsafeNew n :: IO (VUM.IOVector Double)) >>= \c -> VUM.unsafeWrite c 0 (fromMaybe 0 $ readDouble x) >> return c- )- _ ->- MBoxedColumn- <$> ( (VM.unsafeNew n :: IO (VM.IOVector T.Text)) >>= \c -> VM.unsafeWrite c 0 x >> return c- )- VM.unsafeWrite mCol i col-{-# INLINE getInitialDataVectors #-}---- | Reads rows from the handle and stores values in mutable vectors.-fillColumns ::- Int ->- Char ->- VM.IOVector MutableColumn ->- VM.IOVector [(Int, T.Text)] ->- T.Text ->- Handle ->- IO (T.Text, Int)-fillColumns n c mutableCols nullIndices unused handle = do- input <- newIORef unused- rowsRead' <- newIORef (0 :: Int)- forM_ [1 .. (n - 1)] $ \i -> do- atEOF <- hIsEOF handle- input' <- readIORef input- unless (atEOF && input' == mempty) $ do- parseWith (TIO.hGetChunk handle) (parseRow c) input' >>= \case- Fail _unconsumed ctx er -> do- erpos <- hTell handle- fail $- "Failed to parse CSV file around "- <> show erpos- <> " byte; due: "- <> show er- <> "; context: "- <> show ctx- Partial _ -> do- fail "Partial handler is called"- Done (unconsumed :: T.Text) (row :: [T.Text]) -> do- writeIORef input unconsumed- modifyIORef rowsRead' (+ 1)- zipWithM_ (writeValue mutableCols nullIndices i) [0 ..] row- l <- readIORef input- r <- readIORef rowsRead'- pure (l, r)-{-# INLINE fillColumns #-}---- | Writes a value into the appropriate column, resizing the vector if necessary.-writeValue ::- VM.IOVector MutableColumn ->- VM.IOVector [(Int, T.Text)] ->- Int ->- Int ->- T.Text ->- IO ()-writeValue mutableCols nullIndices count colIndex value = do- col <- VM.unsafeRead mutableCols colIndex- res <- writeColumn count value col- let modify val = VM.unsafeModify nullIndices ((count, val) :) colIndex- either modify (const (return ())) res-{-# INLINE writeValue #-}---- | Freezes a mutable vector into an immutable one, trimming it to the actual row count.-freezeColumn ::- V.Vector T.Text ->- VM.IOVector MutableColumn ->- V.Vector [(Int, T.Text)] ->- ReadOptions ->- Int ->- IO Column-freezeColumn colNames mutableCols nulls opts colIndex = do- col <- VM.unsafeRead mutableCols colIndex- let colNulls = nulls V.! colIndex- mode =- effectiveSafeRead- (safeRead opts)- (safeReadOverrides opts)- (colNames V.! colIndex)- case mode of- EitherRead -> freezeColumnEither colNulls col- MaybeRead -> do- frozen <- freezeColumn' colNulls col- return $! ensureOptional frozen- NoSafeRead -> freezeColumn' colNulls col-{-# INLINE freezeColumn #-}---- ------------------------------------------------------------------------------ Streaming scan API--- -----------------------------------------------------------------------------{- | Open a CSV/separated file for streaming, returning an open handle-(positioned just after the header line) and the column specification-for the schema columns that appear in the file header.--The caller is responsible for closing the handle when done.--}-openCsvStream ::- Char ->- Schema ->- FilePath ->- IO (Handle, [(Int, T.Text, SchemaType)])-openCsvStream sep schema path = do- handle <- openFile path ReadMode- hSetBuffering handle (BlockBuffering (Just (8 * 1024 * 1024)))- headerLine <- TIO.hGetLine handle- let headerCols = fmap (T.filter (/= '"') . T.strip) (parseSep sep headerLine)- let schemaMap = elements schema- let colSpec =- [ (idx, name, stype)- | (idx, name) <- zip [0 ..] headerCols- , Just stype <- [M.lookup name schemaMap]- ]- when (null colSpec) $- hClose handle- >> fail- ("openCsvStream: none of the schema columns appear in the header of " <> path)- return (handle, colSpec)--{- | Read up to @batchSz@ rows from the open handle, returning a batch-'DataFrame' and the unconsumed leftover text. Returns 'Nothing' when-the handle is at EOF and there is no leftover input.--The caller must pass the leftover returned by the previous call (use @""@-for the first call).--}-readBatch ::- Char ->- [(Int, T.Text, SchemaType)] ->- Int ->- BS.ByteString ->- Handle ->- IO (Maybe (DataFrame, BS.ByteString))-readBatch sep colSpec batchSz leftover handle = do- let sepByte = fromIntegral (fromEnum sep) :: Word8- numCols = length colSpec- -- Read in 8 MB chunks; only the partial-line tail is copied on refill.- chunkSize = 8 * 1024 * 1024- nullsArr <- VM.unsafeNew numCols- VM.set nullsArr []- mCols <- VM.unsafeNew numCols- forM_ (zip [0 ..] colSpec) $ \(ci, (_, _, st)) ->- VM.unsafeWrite mCols ci =<< makeCol batchSz st- -- buf holds unprocessed bytes; refilled on demand when no newline is found.- bufRef <- newIORef leftover- -- Row-by-row scan. When the buffer has no unquoted newline, fetch another chunk.- -- The copy on refill is only the partial-line tail (≤ one row ≈ few hundred bytes).- let loop !rowIdx = do- remaining <- readIORef bufRef- if rowIdx >= batchSz- then return (rowIdx, remaining)- else case findUnquotedNewline remaining of- Nothing -> do- chunk <- BS.hGet handle chunkSize- if BS.null chunk- then return (rowIdx, remaining) -- EOF- else writeIORef bufRef (remaining <> chunk) >> loop rowIdx- Just nlIdx -> do- let line = BS.take nlIdx remaining- rest' = BS.drop (nlIdx + 1) remaining- line' =- if not (BS.null line) && BS.last line == 0x0D- then BS.init line- else line- writeIORef bufRef rest'- forM_ (zip [0 ..] colSpec) $ \(ci, (fi, _, _)) -> do- let fieldBs = getNthFieldBs sepByte fi line'- col <- VM.unsafeRead mCols ci- res <- writeColumnBs rowIdx fieldBs col- case res of- Left nv -> VM.unsafeModify nullsArr ((rowIdx, nv) :) ci- Right _ -> return ()- loop (rowIdx + 1)- (completeRows, newLeftover) <- loop 0- if completeRows == 0- then return Nothing- else do- forM_ [0 .. numCols - 1] $ \ci -> do- col <- VM.unsafeRead mCols ci- VM.unsafeWrite mCols ci (sliceCol completeRows col)- nullsVec <- V.unsafeFreeze nullsArr- cols <- V.generateM numCols $ \ci -> do- col <- VM.unsafeRead mCols ci- freezeColumn' (nullsVec V.! ci) col- let colNames = [name | (_, name, _) <- colSpec]- return $- Just- ( DataFrame- { columns = cols- , columnIndices = M.fromList (zip colNames [0 ..])- , dataframeDimensions = (completeRows, numCols)- , derivingExpressions = M.empty- }- , newLeftover- )--{- | Write a 'ByteString' field value directly into a mutable column,-parsing numerics without an intermediate 'T.Text' allocation.--}-writeColumnBs ::- Int -> BS.ByteString -> MutableColumn -> IO (Either T.Text Bool)-writeColumnBs i bs (MBoxedColumn (col :: VM.IOVector a)) =- case testEquality (typeRep @a) (typeRep @T.Text) of- Just Refl ->- let val = TextEncoding.decodeUtf8Lenient bs- in VM.unsafeWrite col i val >> return (Right True)- Nothing -> return (Left (TextEncoding.decodeUtf8Lenient bs))-writeColumnBs i bs (MUnboxedColumn (col :: VUM.IOVector a)) =- case testEquality (typeRep @a) (typeRep @Double) of- Just Refl -> case readByteStringDouble bs of- Just v -> VUM.unsafeWrite col i v >> return (Right True)- Nothing -> VUM.unsafeWrite col i 0 >> return (Left (TextEncoding.decodeUtf8Lenient bs))- Nothing -> case testEquality (typeRep @a) (typeRep @Int) of- Just Refl -> case readByteStringInt bs of- Just v -> VUM.unsafeWrite col i v >> return (Right True)- Nothing -> VUM.unsafeWrite col i 0 >> return (Left (TextEncoding.decodeUtf8Lenient bs))- Nothing -> return (Left (TextEncoding.decodeUtf8Lenient bs))-{-# INLINE writeColumnBs #-}--{- | Extracts the Nth field (0-indexed), respecting double quotes and stripping them.-Fast path: uses memchr-based 'BS.break' when no quotes are present in the line.-Slow path: quote-aware character-by-character scan.--}-getNthFieldBs :: Word8 -> Int -> BS.ByteString -> BS.ByteString-getNthFieldBs sep targetIdx bs- | not (BS.any (== 0x22) bs) = skipFast targetIdx bs- | otherwise = go 0 0 False 0- where- -- Fast path: skip fields using elemIndex (memchr); avoids pair allocation.- skipFast k s =- case BS.elemIndex sep s of- Nothing -> if k == 0 then s else BS.empty- Just i ->- if k == 0- then BS.take i s- else skipFast (k - 1) (BS.drop (i + 1) s)-- -- Slow path: quote-aware scan.- quoteChar = 0x22 :: Word8- len = BS.length bs- go !idx !start !inQ !pos- | pos >= len =- if idx == targetIdx then extract start pos else BS.empty- | otherwise =- let c = BS.index bs pos- in if c == quoteChar- then go idx start (not inQ) (pos + 1)- else- if c == sep && not inQ- then- if idx == targetIdx- then extract start pos- else go (idx + 1) (pos + 1) False (pos + 1)- else go idx start inQ (pos + 1)-- extract s e =- let fieldVal = BS.take (e - s) (BS.drop s bs)- in if BS.length fieldVal >= 2- && BS.head fieldVal == quoteChar- && BS.last fieldVal == quoteChar- then BS.init (BS.tail fieldVal)- else fieldVal-{-# INLINE getNthFieldBs #-}---- | Allocate a fresh 'MutableColumn' for @n@ slots based on a 'SchemaType'.-makeCol :: Int -> SchemaType -> IO MutableColumn-makeCol n (SType (_ :: P.Proxy a)) =- case testEquality (typeRep @a) (typeRep @Int) of- Just Refl -> MUnboxedColumn <$> (VUM.unsafeNew n :: IO (VUM.IOVector Int))- Nothing -> case testEquality (typeRep @a) (typeRep @Double) of- Just Refl -> MUnboxedColumn <$> (VUM.unsafeNew n :: IO (VUM.IOVector Double))- Nothing -> MBoxedColumn <$> (VM.unsafeNew n :: IO (VM.IOVector T.Text))---- | Slice a 'MutableColumn' to @n@ elements (no-copy view).-sliceCol :: Int -> MutableColumn -> MutableColumn-sliceCol n (MBoxedColumn col) = MBoxedColumn (VM.take n col)-sliceCol n (MUnboxedColumn col) = MUnboxedColumn (VUM.take n col)--{- | Finds the index of the next unquoted newline (0x0A).-Fast path: uses memchr (SIMD) and falls back to a quote-aware linear scan-only if a double-quote appears before the candidate newline.--}-findUnquotedNewline :: BS.ByteString -> Maybe Int-findUnquotedNewline bs =- case BS.elemIndex 0x0A bs of- Nothing -> Nothing- Just nlPos- -- No quote before the newline → safe to use this position.- -- Check with elemIndex to avoid allocating a ByteString slice.- | maybe True (>= nlPos) (BS.elemIndex 0x22 bs) -> Just nlPos- -- Quote present → may be a newline inside a quoted field; scan carefully.- | otherwise -> slowScan 0 False- where- len = BS.length bs- slowScan !pos !inQ- | pos >= len = Nothing- | otherwise =- let c = BS.index bs pos- in if c == 0x22- then slowScan (pos + 1) (not inQ)- else- if c == 0x0A && not inQ- then Just pos- else slowScan (pos + 1) inQ-{-# INLINE findUnquotedNewline #-}
− src/DataFrame/Lazy/Internal/DataFrame.hs
@@ -1,148 +0,0 @@-{-# LANGUAGE AllowAmbiguousTypes #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE NumericUnderscores #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE ScopedTypeVariables #-}--module DataFrame.Lazy.Internal.DataFrame where--import qualified Data.Text as T-import DataFrame.IO.CSV (CsvReader, readCsvWithSchema)-import qualified DataFrame.Internal.Column as C-import qualified DataFrame.Internal.DataFrame as D-import qualified DataFrame.Internal.Expression as E-import DataFrame.Internal.Schema (Schema)-import DataFrame.Lazy.Internal.Executor (execute)-import DataFrame.Lazy.Internal.LogicalPlan (- DataSource (..),- LogicalPlan (..),- SortOrder (..),- )-import qualified DataFrame.Lazy.Internal.Optimizer as Opt-import DataFrame.Operations.Join (JoinType)--{- | A lazy query that has not been executed yet.--The query is represented as a 'LogicalPlan' tree; execution is deferred-until 'runDataFrame' is called.--}-data LazyDataFrame = LazyDataFrame- { plan :: LogicalPlan- , batchSize :: Int- }--instance Show LazyDataFrame where- show ldf =- "LazyDataFrame { batchSize = "- <> (show (batchSize ldf) <> (", plan = " <> (show (plan ldf) <> " }")))---- ------------------------------------------------------------------------------ Entry point--- -----------------------------------------------------------------------------{- | Execute the lazy query: optimise the logical plan, then stream-execute-the resulting physical plan, returning a fully-materialised 'D.DataFrame'.-The CSV reader (default: attoparsec) is set per scan via 'scanCsv' /-'scanCsvWith'.--}-runDataFrame :: LazyDataFrame -> IO D.DataFrame-runDataFrame ldf = execute (Opt.optimize (batchSize ldf) (plan ldf))---- ------------------------------------------------------------------------------ Builders that construct the logical plan tree--- ------------------------------------------------------------------------------- | Lift an already-loaded eager 'D.DataFrame' into the lazy plan.-fromDataFrame :: D.DataFrame -> LazyDataFrame-fromDataFrame df = LazyDataFrame{plan = SourceDF df, batchSize = 1_000_000}--{- | Scan a CSV file with the default comma separator and the in-tree-attoparsec reader. For the SIMD reader use 'scanCsvWith'.--}-scanCsv :: Schema -> T.Text -> LazyDataFrame-scanCsv = scanCsvWith readCsvWithSchema--{- | Like 'scanCsv' but with an explicit CSV reader (e.g. the SIMD reader-@fastReadCsvWithSchema@ from @dataframe-fastcsv@).--}-scanCsvWith :: CsvReader -> Schema -> T.Text -> LazyDataFrame-scanCsvWith reader schema path =- LazyDataFrame- { plan = Scan (CsvSource (T.unpack path) ',' reader) schema- , batchSize = 1_000_000- }---- | Scan a character-separated file with the default attoparsec reader.-scanSeparated :: Char -> Schema -> T.Text -> LazyDataFrame-scanSeparated = scanSeparatedWith readCsvWithSchema---- | Like 'scanSeparated' but with an explicit CSV reader.-scanSeparatedWith ::- CsvReader -> Char -> Schema -> T.Text -> LazyDataFrame-scanSeparatedWith reader sep schema path =- LazyDataFrame- { plan = Scan (CsvSource (T.unpack path) sep reader) schema- , batchSize = 1_000_000- }---- | Scan a Parquet file, directory of files, or glob pattern.-scanParquet :: Schema -> T.Text -> LazyDataFrame-scanParquet schema path =- LazyDataFrame- { plan = Scan (ParquetSource (T.unpack path)) schema- , batchSize = 1_000_000- }---- | Add a computed column (or overwrite an existing one).-derive ::- (C.Columnable a) => T.Text -> E.Expr a -> LazyDataFrame -> LazyDataFrame-derive name expr ldf =- ldf{plan = Derive name (E.UExpr expr) (plan ldf)}---- | Retain only the listed columns.-select :: [T.Text] -> LazyDataFrame -> LazyDataFrame-select cols ldf = ldf{plan = Project cols (plan ldf)}---- | Keep rows that satisfy the predicate.-filter :: E.Expr Bool -> LazyDataFrame -> LazyDataFrame-filter cond ldf = ldf{plan = Filter cond (plan ldf)}---- | Join two lazy queries on the given key columns.-join ::- JoinType ->- -- | Left join key column name- T.Text ->- -- | Right join key column name- T.Text ->- -- | Left sub-query- LazyDataFrame ->- -- | Right sub-query- LazyDataFrame ->- LazyDataFrame-join jt leftKey rightKey left right =- LazyDataFrame- { plan = Join jt leftKey rightKey (plan left) (plan right)- , batchSize = batchSize left- }--{- | Group by a set of columns and compute aggregate expressions.--Each aggregate expression should use an 'Agg' node (e.g. @sumOf@, @meanOf@).--}-groupBy ::- -- | Group-by key columns- [T.Text] ->- -- | @[(outputName, aggregateExpr)]@- [(T.Text, E.UExpr)] ->- LazyDataFrame ->- LazyDataFrame-groupBy keys aggs ldf = ldf{plan = Aggregate keys aggs (plan ldf)}---- | Sort the result by the given @(column, direction)@ pairs.-sortBy :: [(T.Text, SortOrder)] -> LazyDataFrame -> LazyDataFrame-sortBy cols ldf = ldf{plan = Sort cols (plan ldf)}---- | Retain at most @n@ rows.-take :: Int -> LazyDataFrame -> LazyDataFrame-take n ldf = ldf{plan = Limit n (plan ldf)}
− src/DataFrame/Lazy/Internal/Executor.hs
@@ -1,656 +0,0 @@-{-# LANGUAGE BangPatterns #-}-{-# LANGUAGE ExplicitNamespaces #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TupleSections #-}-{-# LANGUAGE TypeApplications #-}--{- | Pull-based (iterator) execution engine.--Each operator returns a 'Stream' — an IO action that produces the next-'DataFrame' batch on each call and returns 'Nothing' when exhausted.-Blocking operators (Sort, HashJoin) materialise their input before producing-output. HashAggregate uses streaming partial aggregation when all aggregate-expressions support it.--}-module DataFrame.Lazy.Internal.Executor (- CsvReader,- execute,- foldBatches,-) where--import Control.Concurrent (forkIO, getNumCapabilities)-import Control.Concurrent.Async (mapConcurrently)-import Control.Concurrent.STM (atomically)-import Control.Concurrent.STM.TBQueue (newTBQueueIO, readTBQueue, writeTBQueue)-import Control.DeepSeq (force)-import Control.Exception (evaluate)-import Control.Monad (filterM, forM, forM_, when)-import qualified Data.ByteString as BS-import Data.IORef-import qualified Data.Map as M-import qualified Data.Maybe-import qualified Data.Set as S-import qualified Data.Text as T-import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl))-import qualified Data.Vector.Unboxed as VU-import Data.Word (Word8)-import DataFrame.IO.CSV (CsvReader)-import qualified DataFrame.IO.Parquet as Parquet-import qualified DataFrame.Internal.Column as C-import qualified DataFrame.Internal.DataFrame as D-import qualified DataFrame.Internal.Expression as E-import DataFrame.Internal.Schema (elements)-import qualified DataFrame.Lazy.IO.Binary as Bin-import DataFrame.Lazy.Internal.LogicalPlan (DataSource (..), SortOrder (..))-import DataFrame.Lazy.Internal.PhysicalPlan-import qualified DataFrame.Operations.Aggregation as Agg-import qualified DataFrame.Operations.Core as Core-import qualified DataFrame.Operations.Join as Join-import DataFrame.Operations.Merge ()-import qualified DataFrame.Operations.Permutation as Perm-import qualified DataFrame.Operations.Subset as Sub-import qualified DataFrame.Operations.Transformations as Trans-import System.Directory (doesDirectoryExist, removeFile)-import System.FilePath ((</>))-import System.FilePath.Glob (glob)-import System.IO.Temp (emptySystemTempFile)-import Type.Reflection (typeRep)---- ------------------------------------------------------------------------------ Stream abstraction--- -----------------------------------------------------------------------------{- | A pull-based stream: each call to the action yields the next batch or-'Nothing' when the stream is exhausted. State is captured by the closure.--}-newtype Stream = Stream {pullBatch :: IO (Maybe D.DataFrame)}---- | Drain all batches from a stream and concatenate them into one DataFrame.-collectStream :: Stream -> IO D.DataFrame-collectStream stream = go D.empty- where- go acc = do- mb <- pullBatch stream- case mb of- Nothing -> return acc- Just df -> go (acc <> df)---- ------------------------------------------------------------------------------ Top-level entry point--- -----------------------------------------------------------------------------{- | Execute a physical plan, returning the complete result as a single-'DataFrame'.--}-execute :: PhysicalPlan -> IO D.DataFrame-execute plan = buildStream plan >>= collectStream--{- | Fold a function over every batch produced by a physical plan.-The fold is strict in the accumulator; each batch is discarded after folding.--}-foldBatches ::- (b -> D.DataFrame -> IO b) -> b -> PhysicalPlan -> IO b-foldBatches f seed plan = do- stream <- buildStream plan- let loop !acc = do- mb <- pullBatch stream- case mb of- Nothing -> return acc- Just batch -> do- !acc' <- f acc batch- loop acc'- loop seed---- ------------------------------------------------------------------------------ Per-operator stream builders--- -----------------------------------------------------------------------------buildStream :: PhysicalPlan -> IO Stream--- Scan ------------------------------------------------------------------------buildStream (PhysicalScan (CsvSource path sep reader) cfg) =- executeCsvScan path sep reader cfg-buildStream (PhysicalScan (ParquetSource path) cfg) =- executeParquetScan path cfg-buildStream (PhysicalSpill child path) = do- df <- execute child- Bin.spillToDisk path df- df' <- Bin.readSpilled path- ref <- newIORef (Just df')- return . Stream $- ( do- mb <- readIORef ref- writeIORef ref Nothing- return mb- )--- Filter ----------------------------------------------------------------------buildStream (PhysicalFilter p child) = do- childStream <- buildStream child- return . Stream $- ( do- mb <- pullBatch childStream- return $ fmap (Sub.filterWhere p) mb- )--- Project ---------------------------------------------------------------------buildStream (PhysicalProject cols child) = do- childStream <- buildStream child- return . Stream $- ( do- mb <- pullBatch childStream- return $ fmap (Sub.select cols) mb- )--- Derive ----------------------------------------------------------------------buildStream (PhysicalDerive name uexpr child) = do- childStream <- buildStream child- return . Stream $- ( do- mb <- pullBatch childStream- return $ fmap (Trans.deriveMany [(name, uexpr)]) mb- )--- Limit -----------------------------------------------------------------------buildStream (PhysicalLimit n child) = do- childStream <- buildStream child- countRef <- newIORef (0 :: Int)- return . Stream $- ( do- remaining <- readIORef countRef- if remaining >= n- then return Nothing- else do- mb <- pullBatch childStream- case mb of- Nothing -> return Nothing- Just df -> do- let toTake = min (Core.nRows df) (n - remaining)- modifyIORef' countRef (+ toTake)- return $ Just (Sub.take toTake df)- )--- Sort (blocking) -------------------------------------------------------------buildStream (PhysicalSort cols child) = do- df <- execute child- let sortOrds = fmap toPermSortOrder cols- let sorted = Perm.sortBy sortOrds df- ref <- newIORef (Just sorted)- return . Stream $- ( do- mb <- readIORef ref- writeIORef ref Nothing- return mb- )--- HashAggregate ---------------------------------------------------------------buildStream (PhysicalHashAggregate keys aggs child) = do- childStream <- buildStream child- if all (isStreamableAgg . snd) aggs- then do- -- Parallel streaming partial aggregation:- -- * N workers, each pulls batches from the child stream and- -- maintains its own local accumulator.- -- * Once the stream is drained, the N partials are merged- -- sequentially using the same merge expression.- -- * O(|groups| × N) memory in flight, then O(|groups|).- let (partialAggs, mergeAggs, finalizer) = buildAggPlan aggs- nCaps <- getNumCapabilities- let workers = max 1 nCaps- partials <-- mapConcurrently- (\_ -> workerLoop childStream keys partialAggs mergeAggs)- [1 .. workers]- mFinal <-- let nonEmpty = Data.Maybe.catMaybes partials- in case nonEmpty of- [] -> return Nothing- [single] -> return (Just (finalizer single))- (a : rest) -> do- !merged <- mergePartials keys mergeAggs a rest- return (Just (finalizer merged))- ref <- newIORef mFinal- return . Stream $ do- mb <- readIORef ref- writeIORef ref Nothing- return mb- else do- -- Fallback: materialise entire child (for CollectAgg etc.)- df <- collectStream childStream- let result = Agg.aggregate aggs (Agg.groupBy keys df)- ref <- newIORef (Just result)- return . Stream $ do- mb <- readIORef ref- writeIORef ref Nothing- return mb--- SourceDF (split pre-loaded DataFrame into batches) --------------------------buildStream (PhysicalSourceDF bs df) = do- let total = Core.nRows df- posRef <- newIORef (0 :: Int)- return . Stream $ do- i <- readIORef posRef- if i >= total- then return Nothing- else do- let n = min bs (total - i)- batch = Sub.range (i, i + n) df- writeIORef posRef (i + n)- return (Just batch)--- HashJoin — streaming probe (INNER/LEFT) or blocking fallback -----------------buildStream (PhysicalHashJoin jt leftKey rightKey leftPlan rightPlan) =- case jt of- Join.INNER -> streamingHashJoin assembleInnerBatch- Join.LEFT -> streamingHashJoin assembleLeftBatch- _ -> do- -- Blocking fallback for RIGHT / FULL_OUTER- leftDf <- execute leftPlan- rightDf <- execute rightPlan- let result = performJoin jt leftKey rightKey leftDf rightDf- ref <- newIORef (Just result)- return . Stream $ do- mb <- readIORef ref- writeIORef ref Nothing- return mb- where- streamingHashJoin assembleFn = do- -- Materialise build (right) side once and build the compact index.- rightDf <- execute rightPlan- let rightDf' =- if leftKey == rightKey- then rightDf- else Core.rename rightKey leftKey rightDf- joinKey = leftKey- csSet = S.fromList [joinKey]- rightHashes = Join.buildHashColumn [joinKey] rightDf'- ci = Join.buildCompactIndex rightHashes- -- Stream probe (left) side batch by batch.- leftStream <- buildStream leftPlan- return . Stream $ do- mBatch <- pullBatch leftStream- case mBatch of- Nothing -> return Nothing- Just probeBatch -> do- let probeHashes = Join.buildHashColumn [joinKey] probeBatch- (probeIxs, buildIxs) = Join.hashProbeKernel ci probeHashes- return . Just $ assembleFn csSet probeBatch rightDf' probeIxs buildIxs-- assembleLeftBatch csSet probeBatch rightDf' probeIxs buildIxs =- let batchN = Core.nRows probeBatch- -- Mark which probe rows were matched (may have duplicates — that's fine).- matched =- VU.accumulate- (\_ b -> b)- (VU.replicate batchN False)- (VU.map (,True) probeIxs)- unmatchedIxs = VU.findIndices not matched- allProbeIxs = probeIxs VU.++ unmatchedIxs- allBuildIxs = buildIxs VU.++ VU.replicate (VU.length unmatchedIxs) (-1)- in Join.assembleLeft csSet probeBatch rightDf' allProbeIxs allBuildIxs-- assembleInnerBatch = Join.assembleInner---- SortMergeJoin (blocking on both sides) --------------------------------------buildStream (PhysicalSortMergeJoin jt leftKey rightKey leftPlan rightPlan) = do- leftDf <- execute leftPlan- rightDf <- execute rightPlan- let result = performJoin jt leftKey rightKey leftDf rightDf- ref <- newIORef (Just result)- return . Stream $- ( do- mb <- readIORef ref- writeIORef ref Nothing- return mb- )---- ------------------------------------------------------------------------------ Streaming aggregation helpers--- -----------------------------------------------------------------------------{- | True when an aggregate expression can be computed incrementally-(i.e., partial results can be merged without materialising all rows).--}--{- | One worker's loop: pull batches off the shared child stream until-exhausted, building up a per-worker accumulator.--}-workerLoop ::- Stream ->- [T.Text] ->- [E.NamedExpr] ->- [E.NamedExpr] ->- IO (Maybe D.DataFrame)-workerLoop childStream keys partialAggs mergeAggs = loop Nothing- where- loop !acc = do- mb <- pullBatch childStream- case mb of- Nothing -> return acc- Just batch -> do- !partial <-- evaluate . force $- Agg.aggregate partialAggs (Agg.groupBy keys batch)- !next <- case acc of- Nothing -> return (Just partial)- Just a -> do- !merged <-- evaluate . force $- Agg.aggregate mergeAggs (Agg.groupBy keys (a <> partial))- return (Just merged)- loop next---- | Merge a head accumulator with the rest of the workers' partials.-mergePartials ::- [T.Text] ->- [E.NamedExpr] ->- D.DataFrame ->- [D.DataFrame] ->- IO D.DataFrame-mergePartials keys mergeAggs = go- where- go !acc [] = return acc- go !acc (p : ps) = do- !merged <-- evaluate . force $- Agg.aggregate mergeAggs (Agg.groupBy keys (acc <> p))- go merged ps--isStreamableAgg :: E.UExpr -> Bool-isStreamableAgg (E.UExpr (E.Agg (E.CollectAgg _ _) _)) = False-isStreamableAgg (E.UExpr (E.Agg (E.FoldAgg _ Nothing (_ :: a -> b -> a)) _)) =- case testEquality (typeRep @a) (typeRep @b) of- Just Refl -> True -- self-merging: min, max, sum- Nothing -> False-isStreamableAgg (E.UExpr (E.Agg (E.FoldAgg _ (Just _) (_ :: a -> b -> a)) _)) =- case testEquality (typeRep @a) (typeRep @Int) of- Just Refl -> True -- seeded Int fold (old-style count): merge by sum- Nothing ->- case testEquality (typeRep @a) (typeRep @b) of- Just Refl -> True -- seeded self-merging- Nothing -> False-isStreamableAgg (E.UExpr (E.Agg (E.MergeAgg{}) _)) = True-isStreamableAgg _ = False--{- | Build the partial, merge, and finalizer plan for a list of streamable-aggregate expressions.--* @partialAggs@ — applied per batch, producing one row per group-* @mergeAggs@ — applied when combining two partial-result DataFrames-* @finalizer@ — post-process after all batches (needed for 'MergeAgg'- where the accumulator type differs from the output type)--}-buildAggPlan ::- [(T.Text, E.UExpr)] ->- ( [(T.Text, E.UExpr)]- , [(T.Text, E.UExpr)]- , D.DataFrame -> D.DataFrame- )-buildAggPlan aggs = foldl combine ([], [], id) (map processAgg aggs)- where- combine (p1, m1, f1) (p2, m2, f2) = (p1 ++ p2, m1 ++ m2, f1 . f2)-- processAgg ::- (T.Text, E.UExpr) ->- ([(T.Text, E.UExpr)], [(T.Text, E.UExpr)], D.DataFrame -> D.DataFrame)- processAgg (name, ue) = case ue of- -- Seedless FoldAgg: min, max, sum (self-merging when a = b)- E.UExpr (E.Agg (E.FoldAgg n Nothing (f :: a -> b -> a)) (_ :: E.Expr b)) ->- case testEquality (typeRep @a) (typeRep @b) of- Just Refl ->- ( [(name, ue)]- , [(name, E.UExpr (E.Agg (E.FoldAgg n Nothing f) (E.Col @a name)))]- , id- )- Nothing ->- -- a /= b but a = Int: merge by sum (backward compat)- case testEquality (typeRep @a) (typeRep @Int) of- Just Refl ->- ( [(name, ue)]- ,- [- ( name- , E.UExpr- (E.Agg (E.FoldAgg "sum" Nothing ((+) :: Int -> Int -> Int)) (E.Col @Int name))- )- ]- , id- )- Nothing -> ([(name, ue)], [(name, ue)], id)- -- Seeded FoldAgg: old-style count (a = Int)- E.UExpr (E.Agg (E.FoldAgg n (Just _) (f :: a -> b -> a)) (_ :: E.Expr b)) ->- case testEquality (typeRep @a) (typeRep @Int) of- Just Refl ->- ( [(name, ue)]- ,- [- ( name- , E.UExpr- (E.Agg (E.FoldAgg "sum" Nothing ((+) :: Int -> Int -> Int)) (E.Col @Int name))- )- ]- , id- )- Nothing ->- case testEquality (typeRep @a) (typeRep @b) of- Just Refl ->- ( [(name, ue)]- , [(name, E.UExpr (E.Agg (E.FoldAgg n Nothing f) (E.Col @a name)))]- , id- )- Nothing -> ([(name, ue)], [(name, ue)], id)- -- MergeAgg: count, mean, etc.- -- Partial step: accumulate into acc type (using id as finalizer).- -- Merge step: apply merge function to two acc-typed partial results.- -- Finalizer: apply fin to convert acc column to output type.- E.UExpr- ( E.Agg- ( E.MergeAgg- n- seed- (step :: acc -> b -> acc)- (merge :: acc -> acc -> acc)- (fin :: acc -> a)- )- (inner :: E.Expr b)- ) ->- let partialExpr =- E.UExpr- ( E.Agg- (E.MergeAgg n seed step merge (id :: acc -> acc))- inner- )- mergeExpr =- E.UExpr- ( E.Agg- (E.FoldAgg ("merge_" <> n) Nothing merge)- (E.Col @acc name)- )- finalize df =- let accCol = D.unsafeGetColumn name df- finalCol =- either- (error "buildAggPlan: MergeAgg finalize failed")- id- (C.mapColumn @acc @a fin accCol)- in Core.insertColumn name finalCol df- in ( [(name, partialExpr)]- , [(name, mergeExpr)]- , finalize- )- _ -> ([(name, ue)], [(name, ue)], id)---- ------------------------------------------------------------------------------ Parquet scan implementation--- -----------------------------------------------------------------------------{- | Scan a Parquet file, directory, or glob. Each file becomes one batch.-Column projection and predicate pushdown are forwarded to 'readParquetWithOpts'-via 'ParquetReadOptions'.--}-executeParquetScan :: FilePath -> ScanConfig -> IO Stream-executeParquetScan path cfg- | Parquet.isHFUri path = executeHFParquetScan path cfg- | otherwise = do- isDir <- doesDirectoryExist path- let pat = if isDir then path </> "*" else path- matches <- glob pat- files <- filterM (fmap not . doesDirectoryExist) matches- when (null files) $- error ("executeParquetScan: no parquet files found for " ++ path)- let opts =- Parquet.defaultParquetReadOptions- { Parquet.selectedColumns = Just (M.keys (elements (scanSchema cfg)))- , Parquet.predicate = scanPushdownPredicate cfg- }- ref <- newIORef files- return . Stream $ do- fs <- readIORef ref- case fs of- [] -> return Nothing- (f : rest) -> do- writeIORef ref rest- Just <$> Parquet.readParquetWithOpts opts f--{- | HuggingFace Parquet scan. Files are resolved once (API call or direct URL)-then downloaded one at a time as the stream is pulled — so only one file's worth-of data is in memory at a time, regardless of dataset size.--}---- TODO: mchavinda - this should be a more general online file scanner.-executeHFParquetScan :: FilePath -> ScanConfig -> IO Stream-executeHFParquetScan path cfg = do- ref <- case Parquet.parseHFUri path of- Left err -> error err- Right r -> pure r- mToken <- Parquet.getHFToken- hfFiles <-- if Parquet.hasGlob (Parquet.hfGlob ref)- then Parquet.resolveHFUrls mToken ref- else do- let url = Parquet.directHFUrl ref- filename = last $ T.splitOn "/" (Parquet.hfGlob ref)- pure [Parquet.HFParquetFile url "" "" filename]- when (null hfFiles) $- error ("executeParquetScan: no HF parquet files found for " ++ path)- let opts =- Parquet.defaultParquetReadOptions- { Parquet.selectedColumns = Just (M.keys (elements (scanSchema cfg)))- , Parquet.predicate = scanPushdownPredicate cfg- }- filesRef <- newIORef hfFiles- return . Stream $ do- fs <- readIORef filesRef- case fs of- [] -> return Nothing- (f : rest) -> do- writeIORef filesRef rest- -- Download a single file, read it, then return the batch.- [localPath] <- Parquet.downloadHFFiles mToken [f]- Just <$> Parquet.readParquetWithOpts opts localPath---- ------------------------------------------------------------------------------ CSV scan implementation--- -----------------------------------------------------------------------------{- | CSV scan, SIMD-parallel.--The file is read once into memory, split at newline boundaries into N-ByteString slices (N = RTS capabilities), and each slice is parsed in-parallel with the SIMD reader from "DataFrame.IO.CSV.Fast" via the-in-memory entry point — no temp-file roundtrip. The resulting per-chunk-DataFrames are sliced into batches and a dedicated thread feeds them-into a bounded queue. Pushdown predicates are applied per batch by the-consumer.--}-executeCsvScan :: FilePath -> Char -> CsvReader -> ScanConfig -> IO Stream-executeCsvScan path _sep reader cfg = do- nCaps <- getNumCapabilities- chunkPaths <- splitCsvAtNewlines (max 1 nCaps) path-- -- Each chunk parses in parallel via the reader carried on the- -- 'CsvSource' plan node. Parsing and queue-feeding stay disjoint to- -- avoid 14 producers all hammering a shared TBQueue (STM contention- -- dominates throughput).- let schema = scanSchema cfg- batchSz = scanBatchSize cfg- chunkDfs <- mapConcurrently (reader schema) chunkPaths- mapM_ removeFile chunkPaths-- -- Bounded queue with a single writer, N concurrent readers.- queue <- newTBQueueIO (fromIntegral (max 4 (2 * nCaps)))- _ <- forkIO $ do- forM_ chunkDfs $ \df ->- forM_ (sliceIntoBatches batchSz df) $ \b ->- atomically (writeTBQueue queue (Just b))- atomically (writeTBQueue queue Nothing)- return . Stream $- ( do- mb <- atomically (readTBQueue queue)- case mb of- -- Re-insert the sentinel so repeated pulls after EOF stay Nothing.- Nothing -> atomically (writeTBQueue queue Nothing) >> return Nothing- Just df ->- let df' = case scanPushdownPredicate cfg of- Nothing -> df- Just p -> Sub.filterWhere p df- in return (Just df')- )---- | Slice a 'DataFrame' into row-bounded batches of at most @n@ rows.-sliceIntoBatches :: Int -> D.DataFrame -> [D.DataFrame]-sliceIntoBatches n df =- let total = Core.nRows df- starts = [0, n .. total - 1]- in [Sub.range (s, min (s + n) total) df | s <- starts]--{- | Split a CSV file at newline boundaries into @n@ temp files, each-carrying the original header followed by an aligned-at-newlines slice-of the body. Returns the temp file paths; the caller is responsible-for removing them after use. The path-based 'fastReadCsvWithSchema'-mmap's each file, so we get OS-paged reads instead of a single-monolithic 'BS.readFile' of the whole input.--}-splitCsvAtNewlines :: Int -> FilePath -> IO [FilePath]-splitCsvAtNewlines n path = do- bs <- BS.readFile path- let (header, rest) = BS.break (== nl) bs- body = BS.drop 1 rest- bodyLen = BS.length body- rawOffsets = [(bodyLen * i) `div` n | i <- [0 .. n]]- snapped = 0 : map (snap body) (init (drop 1 rawOffsets)) ++ [bodyLen]- ranges = zip snapped (drop 1 snapped)- slices =- [ BS.take (hi - lo) (BS.drop lo body)- | (lo, hi) <- ranges- , hi > lo- ]- forM slices $ \chunk -> do- p <- emptySystemTempFile "lazy_csv_chunk_.csv"- BS.writeFile p (header <> BS.singleton nl <> chunk)- return p- where- nl :: Word8- nl = 0x0A- snap body off =- case BS.elemIndex nl (BS.drop off body) of- Just i -> off + i + 1- Nothing -> BS.length body---- ------------------------------------------------------------------------------ Join helper--- -----------------------------------------------------------------------------{- | Route join to the existing Operations.Join implementation.-When the left and right key names differ, rename the right key before joining.--}-performJoin ::- Join.JoinType -> T.Text -> T.Text -> D.DataFrame -> D.DataFrame -> D.DataFrame-performJoin jt leftKey rightKey leftDf rightDf =- if leftKey == rightKey- then Join.join jt [leftKey] rightDf leftDf- else- let rightRenamed = Core.rename rightKey leftKey rightDf- in Join.join jt [leftKey] rightRenamed leftDf---- ------------------------------------------------------------------------------ Sort order conversion--- ------------------------------------------------------------------------------- | Convert plan-level sort order to the Permutation module's SortOrder.-toPermSortOrder :: (T.Text, SortOrder) -> Perm.SortOrder-toPermSortOrder (col, Ascending) = Perm.Asc (E.Col @T.Text col)-toPermSortOrder (col, Descending) = Perm.Desc (E.Col @T.Text col)
− src/DataFrame/Lazy/Internal/LogicalPlan.hs
@@ -1,49 +0,0 @@-{-# LANGUAGE GADTs #-}--module DataFrame.Lazy.Internal.LogicalPlan where--import qualified Data.Text as T-import DataFrame.IO.CSV (CsvReader)-import qualified DataFrame.Internal.DataFrame as D-import qualified DataFrame.Internal.Expression as E-import DataFrame.Internal.Schema (Schema)-import DataFrame.Operations.Join (JoinType)---- | Data source for a scan node.-data DataSource- = -- | path, separator, CSV reader (e.g. attoparsec or SIMD)- CsvSource FilePath Char CsvReader- | ParquetSource FilePath--instance Show DataSource where- show (CsvSource path sep _) =- "CsvSource " ++ show path ++ " " ++ show sep ++ " <reader>"- show (ParquetSource path) = "ParquetSource " ++ show path---- | Sort direction used in Sort nodes and the public API.-data SortOrder = Ascending | Descending- deriving (Show, Eq, Ord)--{- | Relational-algebra tree that represents what the query computes.-No physical decisions (batch size, join strategy) are made here.--}-data LogicalPlan- = -- | Read columns described by the schema from a source.- Scan DataSource Schema- | -- | Retain only the listed columns.- Project [T.Text] LogicalPlan- | -- | Keep rows matching the predicate.- Filter (E.Expr Bool) LogicalPlan- | -- | Add or overwrite a column via an expression.- Derive T.Text E.UExpr LogicalPlan- | -- | Join two sub-plans on the given key columns.- Join JoinType T.Text T.Text LogicalPlan LogicalPlan- | -- | Group then aggregate.- Aggregate [T.Text] [(T.Text, E.UExpr)] LogicalPlan- | -- | Sort by a list of (column, direction) pairs.- Sort [(T.Text, SortOrder)] LogicalPlan- | -- | Retain at most N rows.- Limit Int LogicalPlan- | -- | Lift an already-loaded DataFrame into the lazy plan.- SourceDF D.DataFrame- deriving (Show)
− src/DataFrame/Lazy/Internal/Optimizer.hs
@@ -1,209 +0,0 @@-{-# LANGUAGE OverloadedStrings #-}--module DataFrame.Lazy.Internal.Optimizer (optimize) where--import qualified Data.Map as M-import qualified Data.Set as S-import qualified Data.Text as T-import qualified DataFrame.Internal.Expression as E-import DataFrame.Internal.Schema (Schema (..), elements)-import DataFrame.Lazy.Internal.LogicalPlan-import DataFrame.Lazy.Internal.PhysicalPlan--{- | Optimise a logical plan and lower it to a physical plan.--Rules applied bottom-up (in order):- 1. Filter fusion — merge consecutive Filter nodes into a conjunction- 2. Predicate pushdown — move Filter past Derive/Project toward Scan- 3. Dead column elim — drop Derive nodes whose output is never referenced--After rule application @toPhysical@ selects concrete operators.--}-optimize :: Int -> LogicalPlan -> PhysicalPlan-optimize batchSz =- toPhysical batchSz- . eliminateDeadColumns- . pushPredicates- . fuseFilters---- ------------------------------------------------------------------------------ Rule 1: Filter fusion--- ------------------------------------------------------------------------------- | Merge @Filter p1 (Filter p2 child)@ into @Filter (p1 && p2) child@.-fuseFilters :: LogicalPlan -> LogicalPlan-fuseFilters (Filter p1 (Filter p2 child)) =- fuseFilters (Filter (andExpr p1 p2) (fuseFilters child))-fuseFilters (Filter p child) = Filter p (fuseFilters child)-fuseFilters (Project cols child) = Project cols (fuseFilters child)-fuseFilters (Derive name expr child) = Derive name expr (fuseFilters child)-fuseFilters (Join jt l r left right) =- Join jt l r (fuseFilters left) (fuseFilters right)-fuseFilters (Aggregate keys aggs child) =- Aggregate keys aggs (fuseFilters child)-fuseFilters (Sort cols child) = Sort cols (fuseFilters child)-fuseFilters (Limit n child) = Limit n (fuseFilters child)-fuseFilters leaf = leaf---- | Logical AND of two @Bool@ expressions.-andExpr :: E.Expr Bool -> E.Expr Bool -> E.Expr Bool-andExpr =- E.Binary- ( E.MkBinaryOp- { E.binaryFn = (&&)- , E.binaryName = "and"- , E.binarySymbol = Just "&&"- , E.binaryCommutative = True- , E.binaryPrecedence = 3- }- )---- ------------------------------------------------------------------------------ Rule 2: Predicate pushdown--- -----------------------------------------------------------------------------{- | Push Filter nodes as close to the Scan as possible.--* Past a @Derive@ when the predicate doesn't reference the derived column.-* Past a @Project@ when all predicate columns are in the projected set.-* Into @ScanConfig.scanPushdownPredicate@ when the child is a @Scan@.--}-pushPredicates :: LogicalPlan -> LogicalPlan-pushPredicates (Filter p (Derive name expr child))- | name `notElem` E.getColumns p =- Derive name expr (pushPredicates (Filter p child))- | otherwise =- Filter p (Derive name expr (pushPredicates child))-pushPredicates (Filter p (Project cols child))- | all (`elem` cols) (E.getColumns p) =- Project cols (pushPredicates (Filter p child))- | otherwise =- Filter p (Project cols (pushPredicates child))-pushPredicates (Filter p child) = Filter p (pushPredicates child)-pushPredicates (Project cols child) = Project cols (pushPredicates child)-pushPredicates (Derive name expr child) = Derive name expr (pushPredicates child)-pushPredicates (Join jt l r left right) =- Join jt l r (pushPredicates left) (pushPredicates right)-pushPredicates (Aggregate keys aggs child) =- Aggregate keys aggs (pushPredicates child)-pushPredicates (Sort cols child) = Sort cols (pushPredicates child)-pushPredicates (Limit n child) = Limit n (pushPredicates child)-pushPredicates leaf = leaf---- ------------------------------------------------------------------------------ Rule 3: Dead column elimination--- -----------------------------------------------------------------------------{- | Collect every column name that is explicitly referenced somewhere in the-plan (in filter predicates, sort keys, aggregate keys, projection lists,-join keys, and derived expressions). Returns Nothing when "all columns-are needed" (i.e. no Project restricts the output).--}-referencedCols :: LogicalPlan -> Maybe (S.Set T.Text)-referencedCols (Scan _ schema) = Just (S.fromList (M.keys (elements schema)))-referencedCols (Project cols _) = Just (S.fromList cols)-referencedCols (Filter p child) =- fmap (S.union (S.fromList (E.getColumns p))) (referencedCols child)-referencedCols (Derive _ expr child) =- fmap (S.union (S.fromList (uExprCols expr))) (referencedCols child)-referencedCols (Join _ l r left right) =- let keySet = S.fromList [l, r]- lRef = fmap (S.union keySet) (referencedCols left)- rRef = fmap (S.union keySet) (referencedCols right)- in liftMaybe2 S.union lRef rRef-referencedCols (Aggregate keys aggs child) =- let aggCols = S.fromList (keys <> concatMap (uExprCols . snd) aggs)- in fmap (S.union aggCols) (referencedCols child)-referencedCols (Sort cols child) =- fmap (S.union (S.fromList (fmap fst cols))) (referencedCols child)-referencedCols (Limit _ child) = referencedCols child-referencedCols (SourceDF _) = Nothing--liftMaybe2 :: (a -> b -> c) -> Maybe a -> Maybe b -> Maybe c-liftMaybe2 f (Just a) (Just b) = Just (f a b)-liftMaybe2 _ _ _ = Nothing--uExprCols :: E.UExpr -> [T.Text]-uExprCols (E.UExpr expr) = E.getColumns expr---- | Drop @Derive@ nodes whose output column is never consumed downstream.-eliminateDeadColumns :: LogicalPlan -> LogicalPlan-eliminateDeadColumns plan = go (referencedCols plan) plan- where- go needed (Derive name expr child) =- case needed of- Nothing -> Derive name expr (go needed child)- Just cols- | name `S.notMember` cols -> go needed child- | otherwise ->- Derive- name- expr- (go (Just (S.union cols (S.fromList (uExprCols expr)))) child)- go needed (Filter p child) =- Filter p (go (fmap (S.union (S.fromList (E.getColumns p))) needed) child)- go _needed (Project cols child) =- Project cols (go (Just (S.fromList cols)) child)- go needed (Join jt l r left right) =- let keySet = fmap (S.union (S.fromList [l, r])) needed- in Join jt l r (go keySet left) (go keySet right)- go needed (Aggregate keys aggs child) =- let aggCols = fmap (S.union (S.fromList (keys <> concatMap (uExprCols . snd) aggs))) needed- in Aggregate keys aggs (go aggCols child)- go needed (Sort cols child) =- Sort cols (go (fmap (S.union (S.fromList (fmap fst cols))) needed) child)- go needed (Limit n child) = Limit n (go needed child)- go needed (Scan ds schema) =- case needed of- Nothing -> Scan ds schema- Just cols ->- Scan ds (Schema (M.filterWithKey (\k _ -> k `S.member` cols) (elements schema)))- go _ (SourceDF df) = SourceDF df---- ------------------------------------------------------------------------------ Logical → Physical lowering--- -----------------------------------------------------------------------------{- | Lower the (already-optimised) logical plan to a physical plan.--Join strategy: always HashJoin (the executor can fall back to SortMerge-at runtime once statistics are available).--}-toPhysical :: Int -> LogicalPlan -> PhysicalPlan--- Special case: Filter directly on a Scan → push into ScanConfig.-toPhysical batchSz (Filter p (Scan (CsvSource path sep reader) schema)) =- PhysicalScan- (CsvSource path sep reader)- (ScanConfig batchSz sep schema (Just p))-toPhysical batchSz (Scan (CsvSource path sep reader) schema) =- PhysicalScan- (CsvSource path sep reader)- (ScanConfig batchSz sep schema Nothing)-toPhysical batchSz (Filter p (Scan (ParquetSource path) schema)) =- PhysicalScan- (ParquetSource path)- (ScanConfig batchSz ',' schema (Just p))-toPhysical batchSz (Scan (ParquetSource path) schema) =- PhysicalScan- (ParquetSource path)- (ScanConfig batchSz ',' schema Nothing)-toPhysical batchSz (Project cols child) =- PhysicalProject cols (toPhysical batchSz child)-toPhysical batchSz (Filter p child) =- PhysicalFilter p (toPhysical batchSz child)-toPhysical batchSz (Derive name expr child) =- PhysicalDerive name expr (toPhysical batchSz child)-toPhysical batchSz (Join jt l r left right) =- PhysicalHashJoin- jt- l- r- (toPhysical batchSz left)- (toPhysical batchSz right)-toPhysical batchSz (Aggregate keys aggs child) =- PhysicalHashAggregate keys aggs (toPhysical batchSz child)-toPhysical batchSz (Sort cols child) =- PhysicalSort cols (toPhysical batchSz child)-toPhysical batchSz (Limit n child) =- PhysicalLimit n (toPhysical batchSz child)-toPhysical batchSz (SourceDF df) = PhysicalSourceDF batchSz df
− src/DataFrame/Lazy/Internal/PhysicalPlan.hs
@@ -1,36 +0,0 @@-module DataFrame.Lazy.Internal.PhysicalPlan where--import qualified Data.Text as T-import qualified DataFrame.Internal.DataFrame as D-import qualified DataFrame.Internal.Expression as E-import DataFrame.Internal.Schema (Schema)-import DataFrame.Lazy.Internal.LogicalPlan (DataSource, SortOrder)-import DataFrame.Operations.Join (JoinType)---- | Scan-level configuration: batch size, separator, optional pushdowns.-data ScanConfig = ScanConfig- { scanBatchSize :: !Int- , scanSeparator :: !Char- , scanSchema :: !Schema- , scanPushdownPredicate :: !(Maybe (E.Expr Bool))- }- deriving (Show)--{- | Physical plan: every node carries enough information for the executor-to allocate resources and choose algorithms without further analysis.--}-data PhysicalPlan- = PhysicalScan DataSource ScanConfig- | PhysicalProject [T.Text] PhysicalPlan- | PhysicalFilter (E.Expr Bool) PhysicalPlan- | PhysicalDerive T.Text E.UExpr PhysicalPlan- | PhysicalHashJoin JoinType T.Text T.Text PhysicalPlan PhysicalPlan- | PhysicalSortMergeJoin JoinType T.Text T.Text PhysicalPlan PhysicalPlan- | PhysicalHashAggregate [T.Text] [(T.Text, E.UExpr)] PhysicalPlan- | PhysicalSort [(T.Text, SortOrder)] PhysicalPlan- | PhysicalLimit Int PhysicalPlan- | -- | Materialize child to a binary file on disk (used for build sides).- PhysicalSpill PhysicalPlan FilePath- | -- | Emit an already-loaded DataFrame as a stream of batches of size @n@.- PhysicalSourceDF Int D.DataFrame- deriving (Show)
− src/DataFrame/Monad.hs
@@ -1,99 +0,0 @@-{-# LANGUAGE ExplicitNamespaces #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE InstanceSigs #-}-{-# LANGUAGE RankNTypes #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TupleSections #-}--module DataFrame.Monad where--import DataFrame (DataFrame)-import qualified DataFrame as D-import DataFrame.Internal.Column (Columnable)-import DataFrame.Internal.Expression (Expr (..))-import DataFrame.Internal.Nullable (BaseType)-import DataFrame.Operations.Transformations (ImputeOp)--import qualified Data.Text as T-import System.Random---- A re-implementation of the state monad.--- `mtl` might be too heavy a dependency just to get--- a single monad instance.-newtype FrameM a = FrameM {runFrameM_ :: DataFrame -> (DataFrame, a)}--instance Functor FrameM where- fmap :: (a -> b) -> FrameM a -> FrameM b- fmap f (FrameM g) = FrameM $ \df ->- let (df', x) = g df- in (df', f x)--instance Applicative FrameM where- pure x = FrameM (,x)- (<*>) :: FrameM (a -> b) -> FrameM a -> FrameM b- FrameM ff <*> FrameM fx = FrameM $ \df ->- let (df1, f) = ff df- (df2, x) = fx df1- in (df2, f x)--instance Monad FrameM where- (>>=) :: FrameM a -> (a -> FrameM b) -> FrameM b- FrameM g >>= f = FrameM $ \df ->- let (df1, x) = g df- FrameM h = f x- in h df1--modifyM :: (DataFrame -> DataFrame) -> FrameM ()-modifyM f = FrameM $ \df -> (f df, ())--inspectM :: (DataFrame -> b) -> FrameM b-inspectM f = FrameM $ \df -> (df, f df)--deriveM :: (Columnable a) => T.Text -> Expr a -> FrameM (Expr a)-deriveM name expr = FrameM $ \df ->- let df' = D.derive name expr df- in (df', Col name)--renameM :: (Columnable a) => Expr a -> T.Text -> FrameM (Expr a)-renameM (Col oldName) newName = FrameM $ \df ->- let df' = D.rename oldName newName df- in (df', Col newName)-renameM expr newName = deriveM newName expr--filterWhereM :: Expr Bool -> FrameM ()-filterWhereM p = modifyM (D.filterWhere p)--sampleM :: (RandomGen g) => g -> Double -> FrameM ()-sampleM pureGen p = modifyM (D.sample pureGen p)--takeM :: Int -> FrameM ()-takeM n = modifyM (D.take n)--filterJustM :: (Columnable a) => Expr (Maybe a) -> FrameM (Expr a)-filterJustM (Col name) = FrameM $ \df ->- let df' = D.filterJust name df- in (df', Col name)-filterJustM expr =- error $ "Cannot filter on compound expression: " ++ show expr--imputeM ::- (ImputeOp a, Columnable (BaseType a)) =>- Expr a ->- BaseType a ->- FrameM (Expr (BaseType a))-imputeM expr@(Col name) value = FrameM $ \df ->- let df' = D.impute expr value df- in (df', Col name)-imputeM expr _ = error $ "Cannot impute on compound expression: " ++ show expr--runFrameM :: DataFrame -> FrameM a -> (a, DataFrame)-runFrameM df (FrameM action) =- let (df', a) = action df- in (a, df')--evalFrameM :: DataFrame -> FrameM a -> a-evalFrameM df m = fst (runFrameM df m)--execFrameM :: DataFrame -> FrameM a -> DataFrame-execFrameM df m = snd (runFrameM df m)
− src/DataFrame/Operations/Aggregation.hs
@@ -1,160 +0,0 @@-{-# LANGUAGE ExplicitNamespaces #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE LambdaCase #-}-{-# LANGUAGE RankNTypes #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE Strict #-}-{-# LANGUAGE TypeApplications #-}--module DataFrame.Operations.Aggregation (- module DataFrame.Operations.Aggregation,- groupBy,- buildRowToGroup,- changingPoints,-) where--import qualified Data.Text as T-import qualified Data.Vector as V-import qualified Data.Vector.Unboxed as VU-import qualified Data.Vector.Unboxed.Mutable as VUM--import Control.Exception (throw)-import Control.Monad-import Control.Monad.ST (runST)-import Data.Hashable-import Data.Type.Equality (TestEquality (..), type (:~:) (Refl))-import DataFrame.Errors-import DataFrame.Internal.Column (- Column (..),- TypedColumn (..),- atIndicesStable,- bitmapTestBit,- )-import DataFrame.Internal.DataFrame (DataFrame (..), GroupedDataFrame (..))-import DataFrame.Internal.Expression-import DataFrame.Internal.Grouping (buildRowToGroup, changingPoints, groupBy)-import DataFrame.Internal.Interpreter-import DataFrame.Internal.Types-import DataFrame.Operations.Core-import DataFrame.Operations.Subset-import Type.Reflection (typeRep)--computeRowHashes :: [Int] -> DataFrame -> VU.Vector Int-computeRowHashes indices df = runST $ do- let n = fst (dimensions df)- mv <- VUM.new n-- let selectedCols = map (columns df V.!) indices-- forM_ selectedCols $ \case- UnboxedColumn _ (v :: VU.Vector a) ->- case testEquality (typeRep @a) (typeRep @Int) of- Just Refl ->- VU.imapM_- ( \i (x :: Int) -> do- h <- VUM.unsafeRead mv i- VUM.unsafeWrite mv i (hashWithSalt h x)- )- v- Nothing ->- case testEquality (typeRep @a) (typeRep @Double) of- Just Refl ->- VU.imapM_- ( \i (d :: Double) -> do- h <- VUM.unsafeRead mv i- VUM.unsafeWrite mv i (hashWithSalt h (doubleToInt d))- )- v- Nothing ->- case sIntegral @a of- STrue ->- VU.imapM_- ( \i d -> do- let x :: Int- x = fromIntegral @a @Int d- h <- VUM.unsafeRead mv i- VUM.unsafeWrite mv i (hashWithSalt h x)- )- v- SFalse ->- case sFloating @a of- STrue ->- VU.imapM_- ( \i d -> do- let x :: Int- x = doubleToInt (realToFrac d :: Double)- h <- VUM.unsafeRead mv i- VUM.unsafeWrite mv i (hashWithSalt h x)- )- v- SFalse ->- VU.imapM_- ( \i d -> do- let x = hash (show d)- h <- VUM.unsafeRead mv i- VUM.unsafeWrite mv i (hashWithSalt h x)- )- v- BoxedColumn bm (v :: V.Vector a) ->- case testEquality (typeRep @a) (typeRep @T.Text) of- Just Refl ->- V.imapM_- ( \i (t :: T.Text) -> do- h <- VUM.unsafeRead mv i- let h' = case bm of- Just bm' | not (bitmapTestBit bm' i) -> hashWithSalt h (0 :: Int)- _ -> hashWithSalt h t- VUM.unsafeWrite mv i h'- )- v- Nothing ->- V.imapM_- ( \i d -> do- let x = case bm of- Just bm' | not (bitmapTestBit bm' i) -> 0 :: Int- _ -> hash (show d)- h <- VUM.unsafeRead mv i- VUM.unsafeWrite mv i (hashWithSalt h x)- )- v-- VU.unsafeFreeze mv- where- doubleToInt :: Double -> Int- doubleToInt = floor . (* 1000)--{- | Aggregate a grouped dataframe using the expressions given.-All ungrouped columns will be dropped.--}-aggregate :: [NamedExpr] -> GroupedDataFrame -> DataFrame-aggregate aggs gdf@(Grouped df groupingColumns valIndices offs _rowToGroup) =- let- df' =- selectIndices- (VU.map (valIndices VU.!) (VU.init offs))- (select groupingColumns df)-- f (name, UExpr (expr :: Expr a)) d =- let- value = case interpretAggregation @a gdf expr of- Left e -> throw e- Right (UnAggregated _) -> throw $ UnaggregatedException (T.pack $ show expr)- Right (Aggregated (TColumn col)) -> col- in- insertColumn name value d- in- fold f aggs df'--selectIndices :: VU.Vector Int -> DataFrame -> DataFrame-selectIndices xs df =- df- { columns = V.map (atIndicesStable xs) (columns df)- , dataframeDimensions = (VU.length xs, V.length (columns df))- }---- | Filter out all non-unique values in a dataframe.-distinct :: DataFrame -> DataFrame-distinct df = selectIndices (VU.map (indices VU.!) (VU.init os)) df- where- (Grouped _ _ indices os _rtg) = groupBy (columnNames df) df
− src/DataFrame/Operations/Core.hs
@@ -1,977 +0,0 @@-{-# LANGUAGE ExplicitNamespaces #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE RankNTypes #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}--module DataFrame.Operations.Core where--import qualified Data.List as L-import qualified Data.Map as M-import qualified Data.Map.Strict as MS-import qualified Data.Text as T-import qualified Data.Vector as V-import qualified Data.Vector.Generic as VG-import qualified Data.Vector.Unboxed as VU--import Control.Exception (throw)-import Data.Bits (popCount)-import Data.Either-import qualified Data.Foldable as Fold-import Data.Function (on, (&))-import Data.Maybe-import Data.Type.Equality (TestEquality (..))-import DataFrame.Errors-import DataFrame.Internal.Column (- Column (..),- Columnable,- TypedColumn (..),- columnLength,- columnTypeString,- expandColumn,- fromList,- fromVector,- toDoubleVector,- toFloatVector,- toIntVector,- toUnboxedVector,- toVector,- )-import DataFrame.Internal.DataFrame (- DataFrame (..),- columnIndices,- derivingExpressions,- empty,- getColumn,- null,- )-import DataFrame.Internal.Expression-import DataFrame.Internal.Interpreter-import DataFrame.Internal.Parsing (isNullish)-import DataFrame.Internal.Row (Any, mkColumnFromRow)-import Type.Reflection-import Prelude hiding (null)--{- | O(1) Get DataFrame dimensions i.e. (rows, columns)--==== __Example__-@->>> :set -XOverloadedStrings->>> import qualified DataFrame as D->>> df = D.fromNamedColumns [("a", D.fromList [1..100]), ("b", D.fromList [1..100]), ("c", D.fromList [1..100])]->>> D.dimensions df--(100, 3)-@--}-dimensions :: DataFrame -> (Int, Int)-dimensions = dataframeDimensions-{-# INLINE dimensions #-}--{- | O(1) Get number of rows in a dataframe.--==== __Example__-@->>> :set -XOverloadedStrings->>> import qualified DataFrame as D->>> df = D.fromNamedColumns [("a", D.fromList [1..100]), ("b", D.fromList [1..100]), ("c", D.fromList [1..100])]->>> D.nRows df-100-@--}-nRows :: DataFrame -> Int-nRows = fst . dataframeDimensions--{- | O(1) Get number of columns in a dataframe.--==== __Example__-@->>> :set -XOverloadedStrings->>> import qualified DataFrame as D->>> df = D.fromNamedColumns [("a", D.fromList [1..100]), ("b", D.fromList [1..100]), ("c", D.fromList [1..100])]->>> D.nColumns df-3-@--}-nColumns :: DataFrame -> Int-nColumns = snd . dataframeDimensions--{- | O(k) Get column names of the DataFrame in order of insertion.--==== __Example__-@->>> :set -XOverloadedStrings->>> import qualified DataFrame as D->>> df = D.fromNamedColumns [("a", D.fromList [1..100]), ("b", D.fromList [1..100]), ("c", D.fromList [1..100])]->>> D.columnNames df--["a", "b", "c"]-@--}-columnNames :: DataFrame -> [T.Text]-columnNames = map fst . L.sortBy (compare `on` snd) . M.toList . columnIndices-{-# INLINE columnNames #-}--{- | Adds a vector to the dataframe. If the vector has less elements than the dataframe and the dataframe is not empty-the vector is converted to type `Maybe a` filled with `Nothing` to match the size of the dataframe. Similarly,-if the vector has more elements than what's currently in the dataframe, the other columns in the dataframe are-change to `Maybe <Type>` and filled with `Nothing`.--==== __Example__-@->>> :set -XOverloadedStrings->>> import qualified DataFrame as D->>> import qualified Data.Vector as V->>> D.insertVector "numbers" (V.fromList [(1 :: Int)..10]) D.empty----------- numbers---------- Int---------- 1- 2- 3- 4- 5- 6- 7- 8- 9- 10--@--}-insertVector ::- forall a.- (Columnable a) =>- -- | Column Name- T.Text ->- -- | Vector to add to column- V.Vector a ->- -- | DataFrame to add column to- DataFrame ->- DataFrame-insertVector name xs = insertColumn name (fromVector xs)-{-# INLINE insertVector #-}--{- | Adds a foldable collection to the dataframe. If the collection has less elements than the-dataframe and the dataframe is not empty-the collection is converted to type `Maybe a` filled with `Nothing` to match the size of the dataframe. Similarly,-if the collection has more elements than what's currently in the dataframe, the other columns in the dataframe are-change to `Maybe <Type>` and filled with `Nothing`.--Be careful not to insert infinite collections with this function as that will crash the program.--==== __Example__-@->>> :set -XOverloadedStrings->>> import qualified DataFrame as D->>> D.insert "numbers" [(1 :: Int)..10] D.empty----------- numbers---------- Int---------- 1- 2- 3- 4- 5- 6- 7- 8- 9- 10--@--}-insert ::- forall a t.- (Columnable a, Foldable t) =>- -- | Column Name- T.Text ->- -- | Sequence to add to dataframe- t a ->- -- | DataFrame to add column to- DataFrame ->- DataFrame-insert name xs = insertColumn name (fromList (Fold.foldr' (:) [] xs)) -- TODO: Do reflection on container type so we can sometimes avoid the list construction.-{-# INLINE insert #-}--{- | Adds a vector to the dataframe and pads it with a default value if it has less elements than the number of rows.--==== __Example__-@->>> :set -XOverloadedStrings->>> import qualified Data.Vector as V->>> import qualified DataFrame as D->>> df = D.fromNamedColumns [("x", D.fromList [(1 :: Int)..10])]->>> D.insertVectorWithDefault 0 "numbers" (V.fromList [(1 :: Int),2,3]) df---------------- x | numbers-----|---------Int | Int-----|---------1 | 1-2 | 2-3 | 3-4 | 0-5 | 0-6 | 0-7 | 0-8 | 0-9 | 0-10 | 0--@--}-insertVectorWithDefault ::- forall a.- (Columnable a) =>- -- | Default Value- a ->- -- | Column name- T.Text ->- -- | Data to add to column- V.Vector a ->- -- | DataFrame to add the column to- DataFrame ->- DataFrame-insertVectorWithDefault defaultValue name xs d =- let (rows, _) = dataframeDimensions d- values = xs V.++ V.replicate (rows - V.length xs) defaultValue- in insertColumn name (fromVector values) d--{- | Adds a list to the dataframe and pads it with a default value if it has less elements than the number of rows.--==== __Example__-@->>> :set -XOverloadedStrings->>> import qualified DataFrame as D->>> df = D.fromNamedColumns [("x", D.fromList [(1 :: Int)..10])]->>> D.insertWithDefault 0 "numbers" [(1 :: Int),2,3] df---------------- x | numbers-----|---------Int | Int-----|---------1 | 1-2 | 2-3 | 3-4 | 0-5 | 0-6 | 0-7 | 0-8 | 0-9 | 0-10 | 0--@--}-insertWithDefault ::- forall a t.- (Columnable a, Foldable t) =>- -- | Default Value- a ->- -- | Column name- T.Text ->- -- | Data to add to column- t a ->- -- | DataFrame to add the column to- DataFrame ->- DataFrame-insertWithDefault defaultValue name xs d =- let (rows, _) = dataframeDimensions d- xs' = Fold.foldr' (:) [] xs- values = xs' ++ replicate (rows - length xs') defaultValue- in insertColumn name (fromList values) d--{- | /O(n)/ Adds an unboxed vector to the dataframe.--Same as insertVector but takes an unboxed vector. If you insert a vector of numbers through insertVector it will either way be converted-into an unboxed vector so this function saves that extra work/conversion.--}-insertUnboxedVector ::- forall a.- (Columnable a, VU.Unbox a) =>- -- | Column Name- T.Text ->- -- | Unboxed vector to add to column- VU.Vector a ->- -- | DataFrame to add the column to- DataFrame ->- DataFrame-insertUnboxedVector name xs = insertColumn name (UnboxedColumn Nothing xs)--{- | /O(n)/ Add a column to the dataframe.--==== __Example__-@->>> :set -XOverloadedStrings->>> import qualified DataFrame as D->>> D.insertColumn "numbers" (D.fromList [(1 :: Int)..10]) D.empty----------- numbers---------- Int---------- 1- 2- 3- 4- 5- 6- 7- 8- 9- 10--@--}-insertColumn ::- -- | Column Name- T.Text ->- -- | Column to add- Column ->- -- | DataFrame to add the column to- DataFrame ->- DataFrame-insertColumn name column d =- let- (r, c) = dataframeDimensions d- n = max (columnLength column) r- exprs = M.delete name (derivingExpressions d)- in- case M.lookup name (columnIndices d) of- Just i ->- DataFrame- (V.map (expandColumn n) (columns d V.// [(i, column)]))- (columnIndices d)- (n, c)- exprs- Nothing ->- DataFrame- (V.map (expandColumn n) (columns d `V.snoc` column))- (M.insert name c (columnIndices d))- (n, c + 1)- exprs--{- | /O(n)/ Clones a column and places it under a new name in the dataframe.--==== __Example__-@->>> :set -XOverloadedStrings->>> import qualified Data.Vector as V->>> df = insertVector "numbers" (V.fromList [1..10]) D.empty->>> D.cloneColumn "numbers" "others" df-------------------- numbers | others----------|-------- Int | Int----------|-------- 1 | 1- 2 | 2- 3 | 3- 4 | 4- 5 | 5- 6 | 6- 7 | 7- 8 | 8- 9 | 9- 10 | 10--@--}-cloneColumn :: T.Text -> T.Text -> DataFrame -> DataFrame-cloneColumn original new df- | null df = throw (EmptyDataSetException "cloneColumn")- | otherwise = fromMaybe- ( throw $- ColumnsNotFoundException [original] "cloneColumn" (M.keys $ columnIndices df)- )- $ do- column <- getColumn original df- return $ insertColumn new column df--{- | /O(n)/ Renames a single column.--==== __Example__-@->>> :set -XOverloadedStrings->>> import qualified DataFrame as D->>> import qualified Data.Vector as V->>> df = insertVector "numbers" (V.fromList [1..10]) D.empty->>> D.rename "numbers" "others" df---------- others--------- Int--------- 1- 2- 3- 4- 5- 6- 7- 8- 9- 10--@--}-rename :: T.Text -> T.Text -> DataFrame -> DataFrame-rename orig new df = either throw id (renameSafe orig new df)--{- | /O(n)/ Renames many columns.--==== __Example__-@->>> :set -XOverloadedStrings->>> import qualified DataFrame as D->>> import qualified Data.Vector as V->>> df = D.insertVector "others" (V.fromList [11..20]) (D.insertVector "numbers" (V.fromList [1..10]) D.empty)->>> df-------------------- numbers | others----------|-------- Int | Int----------|-------- 1 | 11- 2 | 12- 3 | 13- 4 | 14- 5 | 15- 6 | 16- 7 | 17- 8 | 18- 9 | 19- 10 | 20-->>> D.renameMany [("numbers", "first_10"), ("others", "next_10")] df---------------------- first_10 | next_10-----------|--------- Int | Int-----------|--------- 1 | 11- 2 | 12- 3 | 13- 4 | 14- 5 | 15- 6 | 16- 7 | 17- 8 | 18- 9 | 19- 10 | 20--@--}-renameMany :: [(T.Text, T.Text)] -> DataFrame -> DataFrame-renameMany = fold (uncurry rename)--renameSafe ::- T.Text -> T.Text -> DataFrame -> Either DataFrameException DataFrame-renameSafe orig new df- | null df = throw (EmptyDataSetException "rename")- | otherwise = fromMaybe- (Left $ ColumnsNotFoundException [orig] "rename" (M.keys $ columnIndices df))- $ do- columnIndex <- M.lookup orig (columnIndices df)- let origRemoved = M.delete orig (columnIndices df)- let newAdded = M.insert new columnIndex origRemoved- return (Right df{columnIndices = newAdded})--data ColumnInfo = ColumnInfo- { nameOfColumn :: !T.Text- , nonNullValues :: !Int- , nullValues :: !Int- , typeOfColumn :: !T.Text- }--{- | O(n * k ^ 2) Returns the number of non-null columns in the dataframe and the type associated with each column.--==== __Example__-@->>> import qualified Data.Vector as V->>> df = D.insertVector "others" (V.fromList [11..20]) (D.insertVector "numbers" (V.fromList [1..10]) D.empty)->>> D.describeColumns df----------------------------------------------------------- Column Name | # Non-null Values | # Null Values | Type--------------|-------------------|---------------|------ Text | Int | Int | Text--------------|-------------------|---------------|------ others | 10 | 0 | Int- numbers | 10 | 0 | Int--@--}-describeColumns :: DataFrame -> DataFrame-describeColumns df =- empty- & insertColumn "Column Name" (fromList (map nameOfColumn infos))- & insertColumn "# Non-null Values" (fromList (map nonNullValues infos))- & insertColumn "# Null Values" (fromList (map nullValues infos))- & insertColumn "Type" (fromList (map typeOfColumn infos))- where- infos =- L.sortBy (compare `on` nonNullValues) (V.ifoldl' go [] (columns df)) ::- [ColumnInfo]- indexMap = M.fromList (map (\(a, b) -> (b, a)) $ M.toList (columnIndices df))- columnName i = M.lookup i indexMap- go acc i col@(BoxedColumn _bm (_c :: V.Vector a)) =- let- cname = columnName i- countNulls = nulls col- columnType = T.pack $ columnTypeString col- in- if isNothing cname- then acc- else- ColumnInfo- (fromMaybe "" cname)- (columnLength col - countNulls)- countNulls- columnType- : acc- go acc i col@(UnboxedColumn _bm _c) =- let- cname = columnName i- countNulls = nulls col- columnType = T.pack $ columnTypeString col- in- if isNothing cname- then acc- else- ColumnInfo- (fromMaybe "" cname)- (columnLength col - countNulls)- countNulls- columnType- : acc--nulls :: Column -> Int-nulls (BoxedColumn (Just bm) xs) =- -- count null bits in bitmap- let n = VG.length xs- in n - VU.foldl' (\acc b -> acc + popCount b) 0 bm-nulls (BoxedColumn Nothing (xs :: V.Vector a)) = case testEquality (typeRep @a) (typeRep @T.Text) of- Just Refl -> VG.length $ VG.filter isNullish xs- Nothing -> case testEquality (typeRep @a) (typeRep @String) of- Just Refl -> VG.length $ VG.filter (isNullish . T.pack) xs- Nothing -> 0-nulls (UnboxedColumn (Just bm) xs) =- let n = VG.length xs- in n - VU.foldl' (\acc b -> acc + popCount b) 0 bm-nulls _ = 0--{- | Creates a dataframe from a list of tuples with name and column.--==== __Example__-@->>> df = D.fromNamedColumns [("numbers", D.fromList [1..10]), ("others", D.fromList [11..20])]->>> df------------------- numbers | others----------|-------- Int | Int----------|-------- 1 | 11- 2 | 12- 3 | 13- 4 | 14- 5 | 15- 6 | 16- 7 | 17- 8 | 18- 9 | 19- 10 | 20--@--}-fromNamedColumns :: [(T.Text, Column)] -> DataFrame-fromNamedColumns = L.foldl' (\df (name, column) -> insertColumn name column df) empty--{- | Create a dataframe from a list of columns. The column names are "0", "1"... etc.-Useful for quick exploration but you should probably always rename the columns after-or drop the ones you don't want.--==== __Example__-@->>> df = D.fromUnnamedColumns [D.fromList [1..10], D.fromList [11..20]]->>> df------------------- 0 | 1------|----- Int | Int------|----- 1 | 11- 2 | 12- 3 | 13- 4 | 14- 5 | 15- 6 | 16- 7 | 17- 8 | 18- 9 | 19- 10 | 20--@--}-fromUnnamedColumns :: [Column] -> DataFrame-fromUnnamedColumns = fromNamedColumns . zip (map (T.pack . show) [(0 :: Int) ..])--{- | Create a dataframe from a list of column names and rows.--==== __Example__-@->>> df = D.fromRows ["A", "B"] [[D.toAny 1, D.toAny 11], [D.toAny 2, D.toAny 12], [D.toAny 3, D.toAny 13]]-->>> df------------- A | B------|----- Int | Int------|----- 1 | 11- 2 | 12- 3 | 13--@--}-fromRows :: [T.Text] -> [[Any]] -> DataFrame-fromRows names rows =- L.foldl'- (\df i -> insertColumn (names !! i) (mkColumnFromRow i rows) df)- empty- [0 .. length names - 1]--{- | O (k * n) Counts the occurences of each value in a given column.--==== __Example__-@->>> df = D.fromUnnamedColumns [D.fromList [1..10], D.fromList [11..20]]-->>> D.valueCounts @Int "0" df--[(1,1),(2,1),(3,1),(4,1),(5,1),(6,1),(7,1),(8,1),(9,1),(10,1)]--@--}-valueCounts ::- forall a. (Ord a, Columnable a) => Expr a -> DataFrame -> [(a, Int)]-valueCounts expr df- | null df = throw (EmptyDataSetException "valueCounts")- | otherwise = case columnAsVector expr df of- Left e -> throw e- Right column' ->- let- column = V.foldl' (\m v -> MS.insertWith (+) v (1 :: Int) m) M.empty column'- in- M.toAscList column--{- | O (k * n) Shows the proportions of each value in a given column.--==== __Example__-@->>> df = D.fromUnnamedColumns [D.fromList [1..10], D.fromList [11..20]]-->>> D.valueCounts @Int "0" df--[(1,0.1),(2,0.1),(3,0.1),(4,0.1),(5,0.1),(6,0.1),(7,0.1),(8,0.1),(9,0.1),(10,0.1)]--@--}-valueProportions ::- forall a. (Ord a, Columnable a) => Expr a -> DataFrame -> [(a, Double)]-valueProportions expr df- | null df = throw (EmptyDataSetException "valueCounts")- | otherwise = case columnAsVector expr df of- Left e -> throw e- Right column' ->- let- counts =- M.toAscList- (V.foldl' (\m v -> MS.insertWith (+) v (1 :: Int) m) M.empty column')- total = fromIntegral (sum (map snd counts))- in- map (fmap ((/ total) . fromIntegral)) counts--{- | A left fold for dataframes that takes the dataframe as the last object.-This makes it easier to chain operations.--==== __Example__-@->>> df = D.fromNamedColumns [("x", D.fromList [1..100]), ("y", D.fromList [11..110])]->>> D.fold D.dropLast [1..5] df------------ x | y-----|-----Int | Int-----|-----1 | 11-2 | 12-3 | 13-4 | 14-5 | 15-6 | 16-7 | 17-8 | 18-9 | 19-10 | 20-11 | 21-12 | 22-13 | 23-14 | 24-15 | 25-16 | 26-17 | 27-18 | 28-19 | 29-20 | 30--Showing 20 rows out of 85--@--}-fold :: (a -> DataFrame -> DataFrame) -> [a] -> DataFrame -> DataFrame-fold f xs acc = L.foldl' (flip f) acc xs--{- | Returns a dataframe as a two dimensional vector of floats.--Converts all columns in the dataframe to float vectors and transposes them-into a row-major matrix representation.--This is useful for handing data over into ML systems.--Returns 'Left' with an error if any column cannot be converted to floats.--}-toFloatMatrix ::- DataFrame -> Either DataFrameException (V.Vector (VU.Vector Float))-toFloatMatrix df = case V.foldl'- (\acc c -> V.snoc <$> acc <*> toFloatVector c)- (Right V.empty :: Either DataFrameException (V.Vector (VU.Vector Float)))- (columns df) of- Left e -> Left e- Right m ->- pure $- V.generate- (fst (dataframeDimensions df))- ( \i ->- foldl- (\acc j -> acc `VU.snoc` ((m VG.! j) VG.! i))- VU.empty- [0 .. (V.length m - 1)]- )--{- | Returns a dataframe as a two dimensional vector of doubles.--Converts all columns in the dataframe to double vectors and transposes them-into a row-major matrix representation.--This is useful for handing data over into ML systems.--Returns 'Left' with an error if any column cannot be converted to doubles.--}-toDoubleMatrix ::- DataFrame -> Either DataFrameException (V.Vector (VU.Vector Double))-toDoubleMatrix df = case V.foldl'- (\acc c -> V.snoc <$> acc <*> toDoubleVector c)- (Right V.empty :: Either DataFrameException (V.Vector (VU.Vector Double)))- (columns df) of- Left e -> Left e- Right m ->- pure $- V.generate- (fst (dataframeDimensions df))- ( \i ->- foldl- (\acc j -> acc `VU.snoc` ((m VG.! j) VG.! i))- VU.empty- [0 .. (V.length m - 1)]- )--{- | Returns a dataframe as a two dimensional vector of ints.--Converts all columns in the dataframe to int vectors and transposes them-into a row-major matrix representation.--This is useful for handing data over into ML systems.--Returns 'Left' with an error if any column cannot be converted to ints.--}-toIntMatrix :: DataFrame -> Either DataFrameException (V.Vector (VU.Vector Int))-toIntMatrix df = case V.foldl'- (\acc c -> V.snoc <$> acc <*> toIntVector c)- (Right V.empty :: Either DataFrameException (V.Vector (VU.Vector Int)))- (columns df) of- Left e -> Left e- Right m ->- pure $- V.generate- (fst (dataframeDimensions df))- ( \i ->- foldl- (\acc j -> acc `VU.snoc` ((m VG.! j) VG.! i))- VU.empty- [0 .. (V.length m - 1)]- )--{- | Get a specific column as a vector.--You must specify the type via type applications.--==== __Examples__-->>> columnAsVector (F.col @Int "age") df-Right [25, 30, 35, ...]-->>> columnAsVector (F.col @Text "name") df-Right ["Alice", "Bob", "Charlie", ...]--}-columnAsVector ::- forall a.- (Columnable a) => Expr a -> DataFrame -> Either DataFrameException (V.Vector a)-columnAsVector expr df- | null df = throw (EmptyDataSetException "columnAsVector")- | otherwise = case expr of- (Col name) -> case getColumn name df of- Just col -> toVector col- Nothing ->- Left $- ColumnsNotFoundException [name] "columnAsVector" (M.keys $ columnIndices df)- _ -> case interpret df expr of- Left e -> throw e- Right (TColumn col) -> toVector col--{- | Retrieves a column as an unboxed vector of 'Int' values.--Returns 'Left' with a 'DataFrameException' if the column cannot be converted to ints.-This may occur if the column contains non-numeric data or values outside the 'Int' range.--}-columnAsIntVector ::- (Columnable a, Num a) =>- Expr a -> DataFrame -> Either DataFrameException (VU.Vector Int)-columnAsIntVector (Col name) df = case getColumn name df of- Just col -> toIntVector col- Nothing ->- Left $- ColumnsNotFoundException [name] "columnAsIntVector" (M.keys $ columnIndices df)-columnAsIntVector expr df = case interpret df expr of- Left e -> throw e- Right (TColumn col) -> toIntVector col--{- | Retrieves a column as an unboxed vector of 'Double' values.--Returns 'Left' with a 'DataFrameException' if the column cannot be converted to doubles.-This may occur if the column contains non-numeric data.--}-columnAsDoubleVector ::- (Columnable a, Num a) =>- Expr a -> DataFrame -> Either DataFrameException (VU.Vector Double)-columnAsDoubleVector (Col name) df = case getColumn name df of- Just col -> toDoubleVector col- Nothing ->- Left $- ColumnsNotFoundException- [name]- "columnAsDoubleVector"- (M.keys $ columnIndices df)-columnAsDoubleVector expr df = case interpret df expr of- Left e -> throw e- Right (TColumn col) -> toDoubleVector col--{- | Retrieves a column as an unboxed vector of 'Float' values.--Returns 'Left' with a 'DataFrameException' if the column cannot be converted to floats.-This may occur if the column contains non-numeric data.--}-columnAsFloatVector ::- (Columnable a, Num a) =>- Expr a -> DataFrame -> Either DataFrameException (VU.Vector Float)-columnAsFloatVector (Col name) df = case getColumn name df of- Just col -> toFloatVector col- Nothing ->- Left $- ColumnsNotFoundException- [name]- "columnAsFloatVector"- (M.keys $ columnIndices df)-columnAsFloatVector expr df = case interpret df expr of- Left e -> throw e- Right (TColumn col) -> toFloatVector col--columnAsUnboxedVector ::- forall a.- (Columnable a, VU.Unbox a) =>- Expr a -> DataFrame -> Either DataFrameException (VU.Vector a)-columnAsUnboxedVector (Col name) df = case getColumn name df of- Just col -> toUnboxedVector col- Nothing ->- Left $- ColumnsNotFoundException- [name]- "columnAsFloatVector"- (M.keys $ columnIndices df)-columnAsUnboxedVector expr df = case interpret df expr of- Left e -> throw e- Right (TColumn col) -> toUnboxedVector col-{-# SPECIALIZE columnAsUnboxedVector ::- Expr Double -> DataFrame -> Either DataFrameException (VU.Vector Double)- #-}-{-# INLINE columnAsUnboxedVector #-}--{- | Get a specific column as a list.--You must specify the type via type applications.--==== __Examples__-->>> columnAsList @Int "age" df-[25, 30, 35, ...]-->>> columnAsList @Text "name" df-["Alice", "Bob", "Charlie", ...]--==== __Throws__--* 'error' - if the column type doesn't match the requested type--}-columnAsList :: forall a. (Columnable a) => Expr a -> DataFrame -> [a]-columnAsList expr df = either throw V.toList (columnAsVector expr df)--{- | Returns the provenance of all columns in the DataFrame as a list of-@(name, expression)@ pairs. Derived columns show their expression;-raw columns show an identity @col \@type name@ expression.--}---- TODO: mchavinda - Expand out these expressions if possible.-showDerivedExpressions :: DataFrame -> [NamedExpr]-showDerivedExpressions df =- let exprs = derivingExpressions df- names = columnNames df- toNamedExpr name = case M.lookup name exprs of- Just uexpr -> (name, uexpr)- Nothing -> (name, identityUExpr name)- in map toNamedExpr names- where- identityUExpr name = case getColumn name df of- Just (BoxedColumn (Just _) (_ :: V.Vector a)) -> UExpr (Col @(Maybe a) name)- Just (BoxedColumn Nothing (_ :: V.Vector a)) -> UExpr (Col @a name)- Just (UnboxedColumn (Just _) (_ :: VU.Vector a)) -> UExpr (Col @(Maybe a) name)- Just (UnboxedColumn Nothing (_ :: VU.Vector a)) -> UExpr (Col @a name)- Nothing -> error $ "showDerivedExpressions: column not found: " ++ T.unpack name
− src/DataFrame/Operations/Join.hs
@@ -1,1102 +0,0 @@-{-# LANGUAGE BangPatterns #-}-{-# LANGUAGE CPP #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE NumericUnderscores #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE RankNTypes #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}--module DataFrame.Operations.Join where--import Control.Applicative ((<|>))-import Control.Exception (throw)-import Control.Monad (forM_, when)-import Control.Monad.ST (ST, runST)-import qualified Data.HashMap.Strict as HM-#if !MIN_VERSION_base(4,20,0)-import Data.List (foldl')-#endif-import qualified Data.Map.Strict as M-import Data.Maybe (fromMaybe)-import Data.STRef (newSTRef, readSTRef, writeSTRef)-import qualified Data.Set as S-import qualified Data.Text as T-import Data.Type.Equality (TestEquality (..))-import qualified Data.Vector as VB-import qualified Data.Vector.Algorithms.Merge as VA-import qualified Data.Vector.Unboxed as VU-import qualified Data.Vector.Unboxed.Mutable as VUM-import DataFrame.Errors (- DataFrameException (ColumnsNotFoundException),- )-import DataFrame.Internal.Column as D-import DataFrame.Internal.DataFrame as D-import DataFrame.Operations.Aggregation as D-import DataFrame.Operations.Core as D-import Type.Reflection---- | Equivalent to SQL join types.-data JoinType- = INNER- | LEFT- | RIGHT- | FULL_OUTER- deriving (Show)---- | Join two dataframes using SQL join semantics.-join ::- JoinType ->- [T.Text] ->- DataFrame -> -- Right hand side- DataFrame -> -- Left hand side- DataFrame-join INNER xs right = innerJoin xs right-join LEFT xs right = leftJoin xs right-join RIGHT xs right = rightJoin xs right-join FULL_OUTER xs right = fullOuterJoin xs right--{- | Row-count threshold for the build side.-When the build side exceeds this, sort-merge join is used-instead of hash join to avoid L3 cache thrashing.--}-joinStrategyThreshold :: Int-joinStrategyThreshold = 500_000--{- | A compact index mapping hash values to contiguous slices of-original row indices. All indices live in a single unboxed vector;-the HashMap stores @(offset, length)@ into that vector.--}-data CompactIndex = CompactIndex- { ciSortedIndices :: {-# UNPACK #-} !(VU.Vector Int)- , ciOffsets :: !(HM.HashMap Int (Int, Int))- }--{- | Build a compact index from a vector of row hashes.-Sorts @(hash, originalIndex)@ pairs by hash, then scans for-contiguous runs to populate the offset map.--}-buildCompactIndex :: VU.Vector Int -> CompactIndex-buildCompactIndex hashes =- let n = VU.length hashes- (sortedHashes, sortedIndices) = sortWithIndices hashes- !offs = buildOffsets sortedHashes n 0 HM.empty- in CompactIndex sortedIndices offs- where- buildOffsets ::- VU.Vector Int ->- Int ->- Int ->- HM.HashMap Int (Int, Int) ->- HM.HashMap Int (Int, Int)- buildOffsets !sh !n !i !acc- | i >= n = acc- | otherwise =- let !h = sh `VU.unsafeIndex` i- !end = findGroupEnd sh h (i + 1) n- in buildOffsets sh n end (HM.insert h (i, end - i) acc)---- | Find the end of a contiguous run of equal values starting at @j@.-findGroupEnd :: VU.Vector Int -> Int -> Int -> Int -> Int-findGroupEnd !v !h !j !n- | j >= n = j- | v `VU.unsafeIndex` j == h = findGroupEnd v h (j + 1) n- | otherwise = j-{-# INLINE findGroupEnd #-}--{- | Sort a hash vector, returning sorted hashes and corresponding original indices.-Sorts an index array using hash values as the comparison key, avoiding the-intermediate pair vector used by the naive zip-then-sort approach.--}-sortWithIndices :: VU.Vector Int -> (VU.Vector Int, VU.Vector Int)-sortWithIndices hashes = runST $ do- let n = VU.length hashes- mv <- VU.thaw (VU.enumFromN 0 n)- VA.sortBy- (\i j -> compare (hashes `VU.unsafeIndex` i) (hashes `VU.unsafeIndex` j))- mv- sortedIdxs <- VU.unsafeFreeze mv- return (VU.unsafeBackpermute hashes sortedIdxs, sortedIdxs)---- | Write the cross product of two index ranges into mutable vectors.-fillCrossProduct ::- VU.Vector Int ->- VU.Vector Int ->- Int ->- Int ->- Int ->- Int ->- VUM.MVector s Int ->- VUM.MVector s Int ->- Int ->- ST s ()-fillCrossProduct !leftSI !rightSI !lStart !lEnd !rStart !rEnd !lv !rv !pos = goL lStart pos- where- !rLen = rEnd - rStart- goL !li !p- | li >= lEnd = return ()- | otherwise = do- let !lOrigIdx = leftSI `VU.unsafeIndex` li- goR lOrigIdx rStart p- goL (li + 1) (p + rLen)- goR !lOrigIdx !ri !q- | ri >= rEnd = return ()- | otherwise = do- VUM.unsafeWrite lv q lOrigIdx- VUM.unsafeWrite rv q (rightSI `VU.unsafeIndex` ri)- goR lOrigIdx (ri + 1) (q + 1)-{-# INLINE fillCrossProduct #-}---- | Compute key-column indices from the column index map.-keyColIndices :: S.Set T.Text -> DataFrame -> [Int]-keyColIndices csSet df = M.elems $ M.restrictKeys (D.columnIndices df) csSet---- | Validate that all requested join keys exist, then return their indices.-validatedKeyColIndices :: T.Text -> S.Set T.Text -> DataFrame -> [Int]-validatedKeyColIndices callPoint csSet df =- let columnIdxs = D.columnIndices df- missingKeys = S.toAscList (csSet `S.difference` M.keysSet columnIdxs)- in case missingKeys of- [] -> M.elems $ M.restrictKeys columnIdxs csSet- _ -> throw (ColumnsNotFoundException missingKeys callPoint (M.keys columnIdxs))---- ============================================================--- Inner Join--- ============================================================--{- | Performs an inner join on two dataframes using the specified key columns.-Returns only rows where the key values exist in both dataframes.--==== __Example__-@-ghci> df = D.fromNamedColumns [("key", D.fromList ["K0", "K1", "K2", "K3"]), ("A", D.fromList ["A0", "A1", "A2", "A3"])]-ghci> other = D.fromNamedColumns [("key", D.fromList ["K0", "K1", "K2"]), ("B", D.fromList ["B0", "B1", "B2"])]-ghci> D.innerJoin ["key"] df other-------------------- key | A | B-------|-----|----- Text | Text| Text-------|-----|----- K0 | A0 | B0- K1 | A1 | B1- K2 | A2 | B2--@--}-innerJoin :: [T.Text] -> DataFrame -> DataFrame -> DataFrame-innerJoin cs left right- | D.null right || D.null left = D.empty- | otherwise = innerJoinNonEmpty cs left right--innerJoinNonEmpty :: [T.Text] -> DataFrame -> DataFrame -> DataFrame-innerJoinNonEmpty cs left right =- let- csSet = S.fromList cs- leftRows = fst (D.dimensions left)- rightRows = fst (D.dimensions right)-- leftKeyIdxs = validatedKeyColIndices "innerJoin" csSet left- rightKeyIdxs = validatedKeyColIndices "innerJoin" csSet right- leftHashes = D.computeRowHashes leftKeyIdxs left- rightHashes = D.computeRowHashes rightKeyIdxs right-- buildRows = min leftRows rightRows- (leftIxs, rightIxs)- | buildRows > joinStrategyThreshold =- sortMergeInnerKernel leftHashes rightHashes- | rightRows <= leftRows =- -- Build on right (smaller or equal), probe with left- hashInnerKernel leftHashes rightHashes- | otherwise =- -- Build on left (smaller), probe with right, swap result- let (!rIxs, !lIxs) = hashInnerKernel rightHashes leftHashes- in (lIxs, rIxs)- in- assembleInner csSet left right leftIxs rightIxs---- | Compute hashes for the given key column names in a DataFrame.-buildHashColumn :: [T.Text] -> DataFrame -> VU.Vector Int-buildHashColumn keys df =- let csSet = S.fromList keys- keyIdxs = validatedKeyColIndices "buildHashColumn" csSet df- in D.computeRowHashes keyIdxs df--{- | Probe one batch of rows against a pre-built 'CompactIndex'.-Returns @(probeExpandedIxs, buildExpandedIxs)@.-Unlike 'hashInnerKernel', does not build the index (it is pre-built once)-and has no cross-product row guard — the caller controls probe batch size.--}-hashProbeKernel ::- -- | Built once from the full right\/build side.- CompactIndex ->- -- | Probe hashes (one batch).- VU.Vector Int ->- (VU.Vector Int, VU.Vector Int)-hashProbeKernel ci probeHashes =- let ciIxs = ciSortedIndices ci- ciOff = ciOffsets ci- (pFrozen, bFrozen) = runST $ do- let !probeN = VU.length probeHashes- initCap = max 1 (min probeN 1_000_000)-- initPv <- VUM.unsafeNew initCap- initBv <- VUM.unsafeNew initCap- pvRef <- newSTRef initPv- bvRef <- newSTRef initBv- capRef <- newSTRef initCap- posRef <- newSTRef (0 :: Int)-- let ensureCapacity needed = do- cap <- readSTRef capRef- when (needed > cap) $ do- let newCap = max needed (cap * 2)- delta = newCap - cap- pv <- readSTRef pvRef- bv <- readSTRef bvRef- newPv <- VUM.unsafeGrow pv delta- newBv <- VUM.unsafeGrow bv delta- writeSTRef pvRef newPv- writeSTRef bvRef newBv- writeSTRef capRef newCap-- go !i- | i >= probeN = return ()- | otherwise = do- let !h = probeHashes `VU.unsafeIndex` i- case HM.lookup h ciOff of- Nothing -> go (i + 1)- Just (!start, !len) -> do- !p <- readSTRef posRef- ensureCapacity (p + len)- pv <- readSTRef pvRef- bv <- readSTRef bvRef- fillBuild i start len p 0 pv bv- writeSTRef posRef (p + len)- go (i + 1)- fillBuild !probeIdx !start !len !p !j !pv !bv- | j >= len = return ()- | otherwise = do- VUM.unsafeWrite pv (p + j) probeIdx- VUM.unsafeWrite bv (p + j) (ciIxs `VU.unsafeIndex` (start + j))- fillBuild probeIdx start len p (j + 1) pv bv- go 0-- !total <- readSTRef posRef- pv <- readSTRef pvRef- bv <- readSTRef bvRef- (,)- <$> VU.unsafeFreeze (VUM.slice 0 total pv)- <*> VU.unsafeFreeze (VUM.slice 0 total bv)- in (VU.force pFrozen, VU.force bFrozen)--{- | Hash-based inner join kernel.-Builds compact index on @buildHashes@ (second arg), probes with-@probeHashes@ (first arg).-Returns @(probeExpandedIndices, buildExpandedIndices)@.-Uses a dynamically growing output buffer to avoid pre-allocating the full-cross-product size (which can be astronomically large for low-cardinality keys).--}--{- | Maximum number of output rows allowed from a join kernel.-Exceeding this limit indicates a cross-product explosion (e.g. low-cardinality keys).--}-maxJoinOutputRows :: Int-maxJoinOutputRows = 500_000_000--hashInnerKernel ::- VU.Vector Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int)-hashInnerKernel probeHashes buildHashes =- let (pFrozen, bFrozen) = runST $ do- let ci = buildCompactIndex buildHashes- ciIxs = ciSortedIndices ci- ciOff = ciOffsets ci- !probeN = VU.length probeHashes- !buildN = VU.length buildHashes- initCap = max 1 (min (probeN + buildN) 1_000_000)-- initPv <- VUM.unsafeNew initCap- initBv <- VUM.unsafeNew initCap- pvRef <- newSTRef initPv- bvRef <- newSTRef initBv- capRef <- newSTRef initCap- posRef <- newSTRef (0 :: Int)-- let ensureCapacity needed = do- cap <- readSTRef capRef- when (needed > cap) $ do- let newCap = max needed (cap * 2)- delta = newCap - cap- pv <- readSTRef pvRef- bv <- readSTRef bvRef- newPv <- VUM.unsafeGrow pv delta- newBv <- VUM.unsafeGrow bv delta- writeSTRef pvRef newPv- writeSTRef bvRef newBv- writeSTRef capRef newCap-- go !i- | i >= probeN = return ()- | otherwise = do- let !h = probeHashes `VU.unsafeIndex` i- case HM.lookup h ciOff of- Nothing -> go (i + 1)- Just (!start, !len) -> do- !p <- readSTRef posRef- when (p + len > maxJoinOutputRows) $- error $- "Join output would exceed "- ++ show maxJoinOutputRows- ++ " rows (cross-product explosion). "- ++ "Consider filtering or using higher-cardinality join keys or using the lazy API."- ensureCapacity (p + len)- pv <- readSTRef pvRef- bv <- readSTRef bvRef- fillBuild i start len p 0 pv bv- writeSTRef posRef (p + len)- go (i + 1)- fillBuild !probeIdx !start !len !p !j !pv !bv- | j >= len = return ()- | otherwise = do- VUM.unsafeWrite pv (p + j) probeIdx- VUM.unsafeWrite bv (p + j) (ciIxs `VU.unsafeIndex` (start + j))- fillBuild probeIdx start len p (j + 1) pv bv- go 0-- !total <- readSTRef posRef- pv <- readSTRef pvRef- bv <- readSTRef bvRef- (,)- <$> VU.unsafeFreeze (VUM.slice 0 total pv)- <*> VU.unsafeFreeze (VUM.slice 0 total bv)- in -- VU.force copies the slice into a compact array, releasing the oversized- -- backing buffer allocated by the doubling strategy.- (VU.force pFrozen, VU.force bFrozen)--{- | Sort-merge inner join kernel.-Sorts both sides by hash, walks in lockstep.-Returns @(leftExpandedIndices, rightExpandedIndices)@.-Uses a dynamically growing output buffer instead of a two-pass count-then-allocate-strategy, which OOMs when low-cardinality keys produce large cross products.--}-sortMergeInnerKernel ::- VU.Vector Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int)-sortMergeInnerKernel leftHashes rightHashes =- let (lFrozen, rFrozen) = runST $ do- let (leftSH, leftSI) = sortWithIndices leftHashes- (rightSH, rightSI) = sortWithIndices rightHashes- !leftN = VU.length leftHashes- !rightN = VU.length rightHashes- initCap = max 1 (min (leftN + rightN) 1_000_000)-- initLv <- VUM.unsafeNew initCap- initRv <- VUM.unsafeNew initCap- lvRef <- newSTRef initLv- rvRef <- newSTRef initRv- capRef <- newSTRef initCap- posRef <- newSTRef (0 :: Int)-- let ensureCapacity needed = do- cap <- readSTRef capRef- when (needed > cap) $ do- let newCap = max needed (cap * 2)- delta = newCap - cap- lv <- readSTRef lvRef- rv <- readSTRef rvRef- newLv <- VUM.unsafeGrow lv delta- newRv <- VUM.unsafeGrow rv delta- writeSTRef lvRef newLv- writeSTRef rvRef newRv- writeSTRef capRef newCap-- fillGroup !li !lEnd !ri !rEnd = do- let !lLen = lEnd - li- !rLen = rEnd - ri- !groupSize = lLen * rLen- !p <- readSTRef posRef- when (p + groupSize > maxJoinOutputRows) $- error $- "Join output would exceed "- ++ show maxJoinOutputRows- ++ " rows (cross-product explosion with group sizes "- ++ show lLen- ++ " × "- ++ show rLen- ++ "). Consider filtering or using higher-cardinality join keys."- ensureCapacity (p + groupSize)- lv <- readSTRef lvRef- rv <- readSTRef rvRef- let goL !lIdx !pos- | lIdx >= lEnd = return ()- | otherwise = do- let !lOrig = leftSI `VU.unsafeIndex` lIdx- goR lOrig ri pos- goL (lIdx + 1) (pos + rLen)- goR !lOrig !rIdx !pos- | rIdx >= rEnd = return ()- | otherwise = do- VUM.unsafeWrite lv pos lOrig- VUM.unsafeWrite rv pos (rightSI `VU.unsafeIndex` rIdx)- goR lOrig (rIdx + 1) (pos + 1)- goL li p- writeSTRef posRef (p + groupSize)-- fill !li !ri- | li >= leftN || ri >= rightN = return ()- | lh < rh = fill (li + 1) ri- | lh > rh = fill li (ri + 1)- | otherwise = do- let !lEnd = findGroupEnd leftSH lh (li + 1) leftN- !rEnd = findGroupEnd rightSH rh (ri + 1) rightN- fillGroup li lEnd ri rEnd- fill lEnd rEnd- where- !lh = leftSH `VU.unsafeIndex` li- !rh = rightSH `VU.unsafeIndex` ri-- fill 0 0-- !total <- readSTRef posRef- lv <- readSTRef lvRef- rv <- readSTRef rvRef- (,)- <$> VU.unsafeFreeze (VUM.slice 0 total lv)- <*> VU.unsafeFreeze (VUM.slice 0 total rv)- in -- VU.force copies the slice into a compact array, releasing the oversized- -- backing buffer allocated by the doubling strategy.- (VU.force lFrozen, VU.force rFrozen)---- | Assemble the result DataFrame for an inner join from expanded index vectors.-assembleInner ::- S.Set T.Text ->- DataFrame ->- DataFrame ->- VU.Vector Int ->- VU.Vector Int ->- DataFrame-assembleInner csSet left right leftIxs rightIxs =- let !resultLen = VU.length leftIxs- leftColSet = S.fromList (D.columnNames left)- rightColNames = D.columnNames right-- -- Pre-expand every column once- expandedLeftCols = VB.map (D.atIndicesStable leftIxs) (D.columns left)- expandedRightCols = VB.map (D.atIndicesStable rightIxs) (D.columns right)-- getExpandedLeft name = do- idx <- M.lookup name (D.columnIndices left)- return (expandedLeftCols `VB.unsafeIndex` idx)-- getExpandedRight name = do- idx <- M.lookup name (D.columnIndices right)- return (expandedRightCols `VB.unsafeIndex` idx)-- -- Base DataFrame: all left columns, expanded- baseDf =- left- { columns = expandedLeftCols- , dataframeDimensions = (resultLen, snd (D.dataframeDimensions left))- , derivingExpressions = M.empty- }-- insertIfPresent _ Nothing df = df- insertIfPresent name (Just c) df = D.insertColumn name c df- in D.fold- ( \name df ->- if S.member name csSet- then df -- Key column already present from left side- else- if S.member name leftColSet- then -- Overlapping non-key column: merge with These- insertIfPresent- name- (D.mergeColumns <$> getExpandedLeft name <*> getExpandedRight name)- df- else -- Right-only column- insertIfPresent name (getExpandedRight name) df- )- rightColNames- baseDf---- ============================================================--- Left Join--- ============================================================--{- | Performs a left join on two dataframes using the specified key columns.-Returns all rows from the left dataframe, with matching rows from the right dataframe.-Non-matching rows will have Nothing/null values for columns from the right dataframe.--==== __Example__-@-ghci> df = D.fromNamedColumns [("key", D.fromList ["K0", "K1", "K2", "K3"]), ("A", D.fromList ["A0", "A1", "A2", "A3"])]-ghci> other = D.fromNamedColumns [("key", D.fromList ["K0", "K1", "K2"]), ("B", D.fromList ["B0", "B1", "B2"])]-ghci> D.leftJoin ["key"] df other--------------------------- key | A | B-------|-----|----------- Text | Text| Maybe Text-------|-----|----------- K0 | A0 | Just "B0"- K1 | A1 | Just "B1"- K2 | A2 | Just "B2"- K3 | A3 | Nothing--@--}-leftJoin :: [T.Text] -> DataFrame -> DataFrame -> DataFrame-leftJoin = leftJoinWithCallPoint "leftJoin"--leftJoinWithCallPoint ::- T.Text -> [T.Text] -> DataFrame -> DataFrame -> DataFrame-leftJoinWithCallPoint callPoint cs left right- | D.null right || D.nRows right == 0 = left- | D.null left || D.nRows left == 0 = D.empty- | otherwise = leftJoinNonEmpty callPoint cs left right--leftJoinNonEmpty :: T.Text -> [T.Text] -> DataFrame -> DataFrame -> DataFrame-leftJoinNonEmpty callPoint cs left right =- let- csSet = S.fromList cs- rightRows = fst (D.dimensions right)-- leftKeyIdxs = validatedKeyColIndices callPoint csSet left- rightKeyIdxs = validatedKeyColIndices callPoint csSet right- leftHashes = D.computeRowHashes leftKeyIdxs left- rightHashes = D.computeRowHashes rightKeyIdxs right-- -- Right is always the build side for left join- (leftIxs, rightIxs)- | rightRows > joinStrategyThreshold =- sortMergeLeftKernel leftHashes rightHashes- | otherwise =- hashLeftKernel leftHashes rightHashes- in- -- rightIxs uses -1 as sentinel for "no match"- assembleLeft csSet left right leftIxs rightIxs--{- | Hash-based left join kernel.-Returns @(leftExpandedIndices, rightExpandedIndices)@ where-right indices use @-1@ as sentinel for unmatched rows.-Uses a dynamically growing output buffer to avoid pre-allocating the full-cross-product size (which can be astronomically large for low-cardinality keys).--}-hashLeftKernel ::- VU.Vector Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int)-hashLeftKernel leftHashes rightHashes = runST $ do- let ci = buildCompactIndex rightHashes- ciIxs = ciSortedIndices ci- ciOff = ciOffsets ci- !leftN = VU.length leftHashes- !rightN = VU.length rightHashes- initCap = max 1 (min (leftN + rightN) 1_000_000)-- initLv <- VUM.unsafeNew initCap- initRv <- VUM.unsafeNew initCap- lvRef <- newSTRef initLv- rvRef <- newSTRef initRv- capRef <- newSTRef initCap- posRef <- newSTRef (0 :: Int)-- let ensureCapacity needed = do- cap <- readSTRef capRef- when (needed > cap) $ do- let newCap = max needed (cap * 2)- delta = newCap - cap- lv <- readSTRef lvRef- rv <- readSTRef rvRef- newLv <- VUM.unsafeGrow lv delta- newRv <- VUM.unsafeGrow rv delta- writeSTRef lvRef newLv- writeSTRef rvRef newRv- writeSTRef capRef newCap-- go !i- | i >= leftN = return ()- | otherwise = do- let !h = leftHashes `VU.unsafeIndex` i- !p <- readSTRef posRef- case HM.lookup h ciOff of- Nothing -> do- ensureCapacity (p + 1)- lv <- readSTRef lvRef- rv <- readSTRef rvRef- VUM.unsafeWrite lv p i- VUM.unsafeWrite rv p (-1)- writeSTRef posRef (p + 1)- Just (!start, !len) -> do- ensureCapacity (p + len)- lv <- readSTRef lvRef- rv <- readSTRef rvRef- fillBuild i start len p 0 lv rv- writeSTRef posRef (p + len)- go (i + 1)- fillBuild !leftIdx !start !len !p !j !lv !rv- | j >= len = return ()- | otherwise = do- VUM.unsafeWrite lv (p + j) leftIdx- VUM.unsafeWrite rv (p + j) (ciIxs `VU.unsafeIndex` (start + j))- fillBuild leftIdx start len p (j + 1) lv rv- go 0-- !total <- readSTRef posRef- lv <- readSTRef lvRef- rv <- readSTRef rvRef- (,)- <$> VU.unsafeFreeze (VUM.slice 0 total lv)- <*> VU.unsafeFreeze (VUM.slice 0 total rv)--{- | Sort-merge left join kernel.-Returns @(leftExpandedIndices, rightExpandedIndices)@ with @-1@ sentinel.-Uses a dynamically growing output buffer instead of a two-pass count-then-allocate-strategy, which OOMs when low-cardinality keys produce large cross products.--}-sortMergeLeftKernel ::- VU.Vector Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int)-sortMergeLeftKernel leftHashes rightHashes = runST $ do- let (leftSH, leftSI) = sortWithIndices leftHashes- (rightSH, rightSI) = sortWithIndices rightHashes- !leftN = VU.length leftHashes- !rightN = VU.length rightHashes- initCap = max 1 (min (leftN + rightN) 1_000_000)-- initLv <- VUM.unsafeNew initCap- initRv <- VUM.unsafeNew initCap- lvRef <- newSTRef initLv- rvRef <- newSTRef initRv- capRef <- newSTRef initCap- posRef <- newSTRef (0 :: Int)-- let ensureCapacity needed = do- cap <- readSTRef capRef- when (needed > cap) $ do- let newCap = max needed (cap * 2)- delta = newCap - cap- lv <- readSTRef lvRef- rv <- readSTRef rvRef- newLv <- VUM.unsafeGrow lv delta- newRv <- VUM.unsafeGrow rv delta- writeSTRef lvRef newLv- writeSTRef rvRef newRv- writeSTRef capRef newCap-- fillGroup !li !lEnd !ri !rEnd = do- let !lLen = lEnd - li- !rLen = rEnd - ri- !groupSize = lLen * rLen- !p <- readSTRef posRef- ensureCapacity (p + groupSize)- lv <- readSTRef lvRef- rv <- readSTRef rvRef- let goL !lIdx !pos- | lIdx >= lEnd = return ()- | otherwise = do- let !lOrig = leftSI `VU.unsafeIndex` lIdx- goR lOrig ri pos- goL (lIdx + 1) (pos + rLen)- goR !lOrig !rIdx !pos- | rIdx >= rEnd = return ()- | otherwise = do- VUM.unsafeWrite lv pos lOrig- VUM.unsafeWrite rv pos (rightSI `VU.unsafeIndex` rIdx)- goR lOrig (rIdx + 1) (pos + 1)- goL li p- writeSTRef posRef (p + groupSize)-- fill !li !ri- | li >= leftN = return ()- | ri >= rightN = fillRemainingLeft li- | lh < rh = do- !p <- readSTRef posRef- ensureCapacity (p + 1)- lv <- readSTRef lvRef- rv <- readSTRef rvRef- VUM.unsafeWrite lv p (leftSI `VU.unsafeIndex` li)- VUM.unsafeWrite rv p (-1)- writeSTRef posRef (p + 1)- fill (li + 1) ri- | lh > rh = fill li (ri + 1)- | otherwise = do- let !lEnd = findGroupEnd leftSH lh (li + 1) leftN- !rEnd = findGroupEnd rightSH rh (ri + 1) rightN- fillGroup li lEnd ri rEnd- fill lEnd rEnd- where- !lh = leftSH `VU.unsafeIndex` li- !rh = rightSH `VU.unsafeIndex` ri-- fillRemainingLeft !i = do- let !remaining = leftN - i- when (remaining > 0) $ do- !p <- readSTRef posRef- ensureCapacity (p + remaining)- lv <- readSTRef lvRef- rv <- readSTRef rvRef- let go !j- | j >= remaining = return ()- | otherwise = do- VUM.unsafeWrite lv (p + j) (leftSI `VU.unsafeIndex` (i + j))- VUM.unsafeWrite rv (p + j) (-1)- go (j + 1)- go 0- writeSTRef posRef (p + remaining)-- fill 0 0-- !total <- readSTRef posRef- lv <- readSTRef lvRef- rv <- readSTRef rvRef- (,)- <$> VU.unsafeFreeze (VUM.slice 0 total lv)- <*> VU.unsafeFreeze (VUM.slice 0 total rv)--{- | Assemble the result DataFrame for a left join.-Right index vectors use @-1@ sentinel, gathered via 'gatherWithSentinel'.--}-assembleLeft ::- S.Set T.Text ->- DataFrame ->- DataFrame ->- VU.Vector Int ->- VU.Vector Int ->- DataFrame-assembleLeft csSet left right leftIxs rightIxs =- let !resultLen = VU.length leftIxs- leftColSet = S.fromList (D.columnNames left)- rightColNames = D.columnNames right-- expandedLeftCols = VB.map (D.atIndicesStable leftIxs) (D.columns left)- expandedRightCols = VB.map (D.gatherWithSentinel rightIxs) (D.columns right)-- getExpandedLeft name = do- idx <- M.lookup name (D.columnIndices left)- return (expandedLeftCols `VB.unsafeIndex` idx)-- getExpandedRight name = do- idx <- M.lookup name (D.columnIndices right)- return (expandedRightCols `VB.unsafeIndex` idx)-- baseDf =- left- { columns = expandedLeftCols- , dataframeDimensions = (resultLen, snd (D.dataframeDimensions left))- , derivingExpressions = M.empty- }-- insertIfPresent _ Nothing df = df- insertIfPresent name (Just c) df = D.insertColumn name c df- in D.fold- ( \name df ->- if S.member name csSet- then df- else- if S.member name leftColSet- then- insertIfPresent- name- (D.mergeColumns <$> getExpandedLeft name <*> getExpandedRight name)- df- else insertIfPresent name (getExpandedRight name) df- )- rightColNames- baseDf--{- | Performs a right join on two dataframes using the specified key columns.-Returns all rows from the right dataframe, with matching rows from the left dataframe.-Non-matching rows will have Nothing/null values for columns from the left dataframe.--==== __Example__-@-ghci> df = D.fromNamedColumns [("key", D.fromList ["K0", "K1", "K2", "K3"]), ("A", D.fromList ["A0", "A1", "A2", "A3"])]-ghci> other = D.fromNamedColumns [("key", D.fromList ["K0", "K1"]), ("B", D.fromList ["B0", "B1"])]-ghci> D.rightJoin ["key"] df other-------------------- key | A | B-------|-----|----- Text | Text| Text-------|-----|----- K0 | A0 | B0- K1 | A1 | B1--@--}-rightJoin ::- [T.Text] -> DataFrame -> DataFrame -> DataFrame-rightJoin cs left right = leftJoinWithCallPoint "rightJoin" cs right left--fullOuterJoin ::- [T.Text] -> DataFrame -> DataFrame -> DataFrame-fullOuterJoin cs left right- | D.null right || D.nRows right == 0 = left- | D.null left || D.nRows left == 0 = right- | otherwise = fullOuterJoinNonEmpty cs left right--fullOuterJoinNonEmpty :: [T.Text] -> DataFrame -> DataFrame -> DataFrame-fullOuterJoinNonEmpty cs left right =- let- csSet = S.fromList cs- leftRows = fst (D.dimensions left)- rightRows = fst (D.dimensions right)-- leftKeyIdxs = validatedKeyColIndices "fullOuterJoin" csSet left- rightKeyIdxs = validatedKeyColIndices "fullOuterJoin" csSet right- leftHashes = D.computeRowHashes leftKeyIdxs left- rightHashes = D.computeRowHashes rightKeyIdxs right-- -- Both sides can have nulls in full outer- (leftIxs, rightIxs)- | max leftRows rightRows > joinStrategyThreshold =- sortMergeFullOuterKernel leftHashes rightHashes- | otherwise =- hashFullOuterKernel leftHashes rightHashes- in- -- Both index vectors use -1 as sentinel- assembleFullOuter csSet left right leftIxs rightIxs--{- | Hash-based full outer join kernel.-Builds compact indices on both sides.-Returns @(leftExpandedIndices, rightExpandedIndices)@ with @-1@ sentinels.--}-hashFullOuterKernel ::- VU.Vector Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int)-hashFullOuterKernel leftHashes rightHashes = runST $ do- let leftCI = buildCompactIndex leftHashes- rightCI = buildCompactIndex rightHashes- leftOff = ciOffsets leftCI- rightOff = ciOffsets rightCI- leftSI = ciSortedIndices leftCI- rightSI = ciSortedIndices rightCI-- -- Count: matched + left-only + right-only- let leftEntries = HM.toList leftOff- rightEntries = HM.toList rightOff-- !matchedCount =- foldl'- ( \acc (h, (_, ll)) ->- case HM.lookup h rightOff of- Nothing -> acc- Just (_, rl) -> acc + ll * rl- )- 0- leftEntries-- !leftOnlyCount =- foldl'- ( \acc (h, (_, ll)) ->- if HM.member h rightOff then acc else acc + ll- )- 0- leftEntries-- !rightOnlyCount =- foldl'- ( \acc (h, (_, rl)) ->- if HM.member h leftOff then acc else acc + rl- )- 0- rightEntries-- !totalCount = matchedCount + leftOnlyCount + rightOnlyCount-- lv <- VUM.unsafeNew totalCount- rv <- VUM.unsafeNew totalCount- posRef <- newSTRef (0 :: Int)-- -- Fill matched + left-only (iterate left keys)- forM_ leftEntries $ \(h, (lStart, lLen)) -> do- !p <- readSTRef posRef- case HM.lookup h rightOff of- Nothing -> do- -- Left-only rows- let goL !j !q- | j >= lLen = return ()- | otherwise = do- VUM.unsafeWrite lv q (leftSI `VU.unsafeIndex` (lStart + j))- VUM.unsafeWrite rv q (-1)- goL (j + 1) (q + 1)- goL 0 p- writeSTRef posRef (p + lLen)- Just (!rStart, !rLen) -> do- -- Cross product- fillCrossProduct- leftSI- rightSI- lStart- (lStart + lLen)- rStart- (rStart + rLen)- lv- rv- p- writeSTRef posRef (p + lLen * rLen)-- -- Fill right-only (iterate right keys not in left)- forM_ rightEntries $ \(h, (rStart, rLen)) ->- case HM.lookup h leftOff of- Just _ -> return ()- Nothing -> do- !p <- readSTRef posRef- let goR !j !q- | j >= rLen = return ()- | otherwise = do- VUM.unsafeWrite lv q (-1)- VUM.unsafeWrite rv q (rightSI `VU.unsafeIndex` (rStart + j))- goR (j + 1) (q + 1)- goR 0 p- writeSTRef posRef (p + rLen)-- (,) <$> VU.unsafeFreeze lv <*> VU.unsafeFreeze rv--{- | Sort-merge full outer join kernel.-Returns @(leftExpandedIndices, rightExpandedIndices)@ with @-1@ sentinels.--}-sortMergeFullOuterKernel ::- VU.Vector Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int)-sortMergeFullOuterKernel leftHashes rightHashes = runST $ do- let (leftSH, leftSI) = sortWithIndices leftHashes- (rightSH, rightSI) = sortWithIndices rightHashes- !leftN = VU.length leftHashes- !rightN = VU.length rightHashes-- -- Pass 1: count- let countLoop !li !ri !c- | li >= leftN && ri >= rightN = c- | li >= leftN = c + (rightN - ri)- | ri >= rightN = c + (leftN - li)- | lh < rh = countLoop (li + 1) ri (c + 1)- | lh > rh = countLoop li (ri + 1) (c + 1)- | otherwise =- let !lEnd = findGroupEnd leftSH lh (li + 1) leftN- !rEnd = findGroupEnd rightSH rh (ri + 1) rightN- in countLoop lEnd rEnd (c + (lEnd - li) * (rEnd - ri))- where- !lh = leftSH `VU.unsafeIndex` li- !rh = rightSH `VU.unsafeIndex` ri- !totalRows = countLoop 0 0 0-- -- Pass 2: fill- lv <- VUM.unsafeNew totalRows- rv <- VUM.unsafeNew totalRows-- let fill !li !ri !pos- | li >= leftN && ri >= rightN = return ()- | li >= leftN = fillRemainingRight ri pos- | ri >= rightN = fillRemainingLeft li pos- | lh < rh = do- VUM.unsafeWrite lv pos (leftSI `VU.unsafeIndex` li)- VUM.unsafeWrite rv pos (-1)- fill (li + 1) ri (pos + 1)- | lh > rh = do- VUM.unsafeWrite lv pos (-1)- VUM.unsafeWrite rv pos (rightSI `VU.unsafeIndex` ri)- fill li (ri + 1) (pos + 1)- | otherwise = do- let !lEnd = findGroupEnd leftSH lh (li + 1) leftN- !rEnd = findGroupEnd rightSH rh (ri + 1) rightN- !groupSize = (lEnd - li) * (rEnd - ri)- fillCrossProduct leftSI rightSI li lEnd ri rEnd lv rv pos- fill lEnd rEnd (pos + groupSize)- where- !lh = leftSH `VU.unsafeIndex` li- !rh = rightSH `VU.unsafeIndex` ri-- fillRemainingLeft !i !pos- | i >= leftN = return ()- | otherwise = do- VUM.unsafeWrite lv pos (leftSI `VU.unsafeIndex` i)- VUM.unsafeWrite rv pos (-1)- fillRemainingLeft (i + 1) (pos + 1)-- fillRemainingRight !i !pos- | i >= rightN = return ()- | otherwise = do- VUM.unsafeWrite lv pos (-1)- VUM.unsafeWrite rv pos (rightSI `VU.unsafeIndex` i)- fillRemainingRight (i + 1) (pos + 1)-- fill 0 0 0- (,) <$> VU.unsafeFreeze lv <*> VU.unsafeFreeze rv--{- | Assemble the result DataFrame for a full outer join.-Both index vectors use @-1@ sentinel; all columns gathered via-'gatherWithSentinel'. Key columns are coalesced (first non-null wins).--}-assembleFullOuter ::- S.Set T.Text ->- DataFrame ->- DataFrame ->- VU.Vector Int ->- VU.Vector Int ->- DataFrame-assembleFullOuter csSet left right leftIxs rightIxs =- let !resultLen = VU.length leftIxs- leftColSet = S.fromList (D.columnNames left)- rightColNames = D.columnNames right-- expandedLeftCols = VB.map (D.gatherWithSentinel leftIxs) (D.columns left)- expandedRightCols = VB.map (D.gatherWithSentinel rightIxs) (D.columns right)-- getExpandedLeft name = do- idx <- M.lookup name (D.columnIndices left)- return (expandedLeftCols `VB.unsafeIndex` idx)-- getExpandedRight name = do- idx <- M.lookup name (D.columnIndices right)- return (expandedRightCols `VB.unsafeIndex` idx)-- baseDf =- left- { columns = expandedLeftCols- , dataframeDimensions = (resultLen, snd (D.dataframeDimensions left))- , derivingExpressions = M.empty- }-- insertIfPresent _ Nothing df = df- insertIfPresent name (Just c) df = D.insertColumn name c df-- -- Coalesce two nullable columns: take first non-Nothing per row,- -- producing a non-optional column.- coalesceKeyColumn :: Column -> Column -> Column- coalesceKeyColumn- (BoxedColumn lBm (lCol :: VB.Vector a))- (BoxedColumn rBm (rCol :: VB.Vector b)) =- case testEquality (typeRep @a) (typeRep @b) of- Just Refl ->- let asMaybe bm =- VB.imap- ( \i v -> case bm of- Just bm' -> if bitmapTestBit bm' i then Just v else Nothing- Nothing -> Just v- )- lMaybe = asMaybe lBm lCol- rMaybe = asMaybe rBm rCol- in D.fromVector $- VB.zipWith- ( \l r ->- fromMaybe (error "fullOuterJoin: null on both sides of key column") (l <|> r)- )- lMaybe- rMaybe- Nothing -> error "Cannot join columns of different types"- coalesceKeyColumn _ _ = error "fullOuterJoin: expected nullable column for key columns"- in D.fold- ( \name df ->- if S.member name csSet- then -- Key column: coalesce left and right- case (getExpandedLeft name, getExpandedRight name) of- (Just lc, Just rc) -> D.insertColumn name (coalesceKeyColumn lc rc) df- _ -> df- else- if S.member name leftColSet- then- insertIfPresent- name- (D.mergeColumns <$> getExpandedLeft name <*> getExpandedRight name)- df- else insertIfPresent name (getExpandedRight name) df- )- rightColNames- baseDf
− src/DataFrame/Operations/Merge.hs
@@ -1,73 +0,0 @@-{-# LANGUAGE InstanceSigs #-}-{-# OPTIONS_GHC -Wno-orphans #-}--module DataFrame.Operations.Merge where--import qualified Data.List as L-import qualified Data.Text as T-import qualified DataFrame.Internal.Column as D-import qualified DataFrame.Internal.DataFrame as D-import qualified DataFrame.Operations.Core as D--import Data.Maybe--{- | Vertically merge two dataframes using shared columns.-Columns that exist in only one dataframe are padded with Nothing.--}-instance Semigroup D.DataFrame where- (<>) :: D.DataFrame -> D.DataFrame -> D.DataFrame- (<>) a b =- let- addColumns a' b' df name- | snd (D.dimensions a') == 0 && snd (D.dimensions b') == 0 = df- | snd (D.dimensions a') == 0 = fromMaybe df $ do- col <- D.getColumn name b'- pure $ D.insertColumn name col df- | snd (D.dimensions b') == 0 = fromMaybe df $ do- col <- D.getColumn name a'- pure $ D.insertColumn name col df- | otherwise =- let- numRowsA = fst $ D.dimensions a'- numRowsB = fst $ D.dimensions b'- sumRows = numRowsA + numRowsB-- optA = D.getColumn name a'- optB = D.getColumn name b'- in- case optB of- Nothing -> case optA of- Nothing ->- -- N.B. this case should never happen, because we're dealing with columns coming from- -- union of column names of both dataframes. Nothing + Nothing would mean column- -- wasn't in either dataframe, which shouldn't happen- D.insertColumn name (D.fromList ([] :: [T.Text])) df- Just a'' ->- D.insertColumn name (D.expandColumn sumRows a'') df- Just b'' -> case optA of- Nothing ->- D.insertColumn name (D.leftExpandColumn sumRows b'') df- Just a'' ->- let concatedColumns = D.concatColumnsEither a'' b''- in D.insertColumn name concatedColumns df- result = L.foldl' (addColumns a b) D.empty (D.columnNames a `L.union` D.columnNames b)- in- result- { D.derivingExpressions = D.derivingExpressions a <> D.derivingExpressions b- }--instance Monoid D.DataFrame where- mempty = D.empty---- | Add two dataframes side by side/horizontally.-(|||) :: D.DataFrame -> D.DataFrame -> D.DataFrame-(|||) a b =- let result =- D.fold- (\name acc -> D.insertColumn name (D.unsafeGetColumn name b) acc)- (D.columnNames b)- a- in result- { D.derivingExpressions =- D.derivingExpressions result <> D.derivingExpressions b- }
− src/DataFrame/Operations/Permutation.hs
@@ -1,164 +0,0 @@-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE InstanceSigs #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}--module DataFrame.Operations.Permutation where--import qualified Data.List as L-import qualified Data.Text as T-import qualified Data.Vector as V-import qualified Data.Vector.Algorithms.Merge as VA-import qualified Data.Vector.Generic as VG-import qualified Data.Vector.Unboxed as VU-import qualified Data.Vector.Unboxed.Mutable as VUM--import Control.Exception (throw)-import Control.Monad.ST (runST)-import Data.Type.Equality (testEquality, (:~:) (Refl))-import Data.Vector.Internal.Check (HasCallStack)-import DataFrame.Errors (DataFrameException (..))-import DataFrame.Internal.Column (Column (..), Columnable, atIndicesStable)-import DataFrame.Internal.DataFrame (DataFrame (..), unsafeGetColumn)-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)---- | Sort order taken as a parameter by the 'sortBy' function.-data SortOrder where- Asc :: (Columnable a, Ord a) => Expr a -> SortOrder- Desc :: (Columnable a, Ord a) => Expr a -> SortOrder--instance Eq SortOrder where- (==) :: SortOrder -> SortOrder -> Bool- (==) (Asc _) (Asc _) = True- (==) (Desc _) (Desc _) = True- (==) _ _ = False--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--}-sortBy ::- [SortOrder] ->- DataFrame ->- DataFrame-sortBy sortOrds df- | not (null missing) =- throw $- ColumnsNotFoundException- missing- "sortBy"- (columnNames df)- | otherwise =- let- (sortOrds', df') = prepareSortColumns sortOrds df- comparators = map (`sortOrderComparator` df') sortOrds'- compositeCompare i j = mconcat [c i j | c <- comparators]- nRows = fst (dataframeDimensions df')- indexes = sortIndices compositeCompare nRows- in- df{columns = V.map (atIndicesStable indexes) (columns df)}- where- 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.--}-sortOrderComparator :: SortOrder -> DataFrame -> Int -> Int -> Ordering-sortOrderComparator (Asc (Col name :: Expr a)) df =- case unsafeGetColumn name df of- BoxedColumn _ (v :: V.Vector b) -> case testEquality (typeRep @a) (typeRep @b) of- Just Refl -> \i j -> compare (v `V.unsafeIndex` i) (v `V.unsafeIndex` j)- Nothing -> \_ _ -> EQ- UnboxedColumn _ (v :: VU.Vector b) -> case testEquality (typeRep @a) (typeRep @b) of- Just Refl -> \i j -> compare (v `VU.unsafeIndex` i) (v `VU.unsafeIndex` j)- Nothing -> \_ _ -> EQ-sortOrderComparator (Desc (Col name :: Expr a)) df =- case unsafeGetColumn name df of- BoxedColumn _ (v :: V.Vector b) -> case testEquality (typeRep @a) (typeRep @b) of- Just Refl -> \i j -> compare (v `V.unsafeIndex` j) (v `V.unsafeIndex` i)- Nothing -> \_ _ -> EQ- UnboxedColumn _ (v :: VU.Vector b) -> case testEquality (typeRep @a) (typeRep @b) of- Just Refl -> \i j -> compare (v `VU.unsafeIndex` j) (v `VU.unsafeIndex` i)- Nothing -> \_ _ -> EQ-sortOrderComparator _ _ = error "Sorting on compound column"---- | Sort row indices using a comparator function.-sortIndices :: (Int -> Int -> Ordering) -> Int -> VU.Vector Int-sortIndices cmp nRows = runST $ do- withIndexes <- VG.thaw (V.generate nRows id :: V.Vector Int)- VA.sortBy cmp withIndexes- sorted <- VG.unsafeFreeze withIndexes- return (VU.convert sorted)--shuffle ::- (RandomGen g) =>- g ->- DataFrame ->- DataFrame-shuffle pureGen df =- let- indexes = shuffledIndices pureGen (fst (dimensions df))- in- df{columns = V.map (atIndicesStable indexes) (columns df)}--shuffledIndices :: (HasCallStack, RandomGen g) => g -> Int -> VU.Vector Int-shuffledIndices pureGen k- | k < 0 = error $ "Vector index may not be a neative number: " <> show k- | k == 0 = VU.empty- | otherwise = shuffleVec pureGen- where- shuffleVec :: (RandomGen g) => g -> VU.Vector Int- shuffleVec g = runST $ do- vm <- VUM.generate k id- let (n, nGen) = randomR (1, k - 1) g- go vm n nGen- VU.unsafeFreeze vm-- go _v (-1) _ = pure ()- go _v 0 _ = pure ()- go v maxInd gen =- let- (n, nextGen) = randomR (1, maxInd) gen- in- VUM.swap v 0 n *> go (VUM.tail v) (maxInd - 1) nextGen
− src/DataFrame/Operations/Statistics.hs
@@ -1,397 +0,0 @@-{-# LANGUAGE ExplicitNamespaces #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE FlexibleInstances #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE RankNTypes #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}-{-# LANGUAGE UndecidableInstances #-}-{-# OPTIONS_GHC -Wno-orphans #-}--module DataFrame.Operations.Statistics where--import qualified Data.List as L-import qualified Data.Map as M-import qualified Data.Text as T-import qualified Data.Vector as V-import qualified Data.Vector.Generic as VG-import qualified Data.Vector.Unboxed as VU--import Prelude as P--import Control.Exception (throw)-import Data.Function ((&))-import Data.Maybe (fromMaybe, isJust)-import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl))-import DataFrame.Errors (DataFrameException (..))-import DataFrame.Internal.Column-import DataFrame.Internal.DataFrame (- DataFrame (..),- empty,- getColumn,- )-import DataFrame.Internal.Expression-import DataFrame.Internal.Interpreter-import DataFrame.Internal.Nullable (BaseType)-import DataFrame.Internal.Row (showValue, toAny)-import DataFrame.Internal.Statistics-import DataFrame.Internal.Types-import DataFrame.Operations.Core-import DataFrame.Operations.Subset (filterJust)-import DataFrame.Operations.Transformations (ImputeOp (..), imputeCore)-import Text.Printf (printf)-import Type.Reflection (typeRep)--{- | Show a frequency table for a categorical feaure.--__Examples:__--@-ghci> df <- D.readCsv ".\/data\/housing.csv"--ghci> D.frequencies "ocean_proximity" df------------------------------------------------------------------------ Statistic | <1H OCEAN | INLAND | ISLAND | NEAR BAY | NEAR OCEAN-----------------|-----------|--------|--------|----------|------------ Text | Any | Any | Any | Any | Any-----------------|-----------|--------|--------|----------|------------ Count | 9136 | 6551 | 5 | 2290 | 2658- Percentage (%) | 44.26% | 31.74% | 0.02% | 11.09% | 12.88%-@--}-frequencies ::- forall a. (Columnable a, Ord a) => Expr a -> DataFrame -> DataFrame-frequencies expr df =- let- counts = valueCounts expr df- calculatePercentage cs k = toAny $ toPct2dp (fromIntegral k / fromIntegral (P.sum $ map snd cs))- initDf =- empty- & insertVector "Statistic" (V.fromList ["Count" :: T.Text, "Percentage (%)"])- freqs _col' =- L.foldl'- ( \d (col'', k) ->- insertVector- (showValue @a col'')- (V.fromList [toAny k, calculatePercentage counts k])- d- )- initDf- counts- in- case columnAsVector expr df of- Left err -> throw err- Right column -> freqs column---- | Calculates the mean of a given column as a standalone value.-mean ::- forall a. (Columnable a, Real a, VU.Unbox a) => Expr a -> DataFrame -> Double-mean (Col name) df = case _getColumnAsDouble name df of- Just xs -> meanDouble' xs- Nothing -> error "[INTERNAL ERROR] Column is non-numeric"-mean expr df = case interpret df expr of- Left e -> throw e- Right (TColumn col) -> case toUnboxedVector @a col of- Left e -> throw e- Right xs -> mean' xs--meanMaybe ::- forall a. (Columnable a, Real a) => Expr (Maybe a) -> DataFrame -> Double-meanMaybe (Col name) df =- (mean' . optionalToDoubleVector)- (either throw id (columnAsVector (Col @(Maybe a) name) df))-meanMaybe expr df = case interpret @(Maybe a) df expr of- Left e -> throw e- Right (TColumn col) -> case toVector @(Maybe a) col of- Left e -> throw e- Right xs -> (mean' . optionalToDoubleVector) xs---- | Calculates the median of a given column as a standalone value.-median ::- forall a. (Columnable a, Real a, VU.Unbox a) => Expr a -> DataFrame -> Double-median (Col name) df = case columnAsUnboxedVector (Col @a name) df of- Right xs -> median' xs- Left e -> throw e-median expr df = case interpret df expr of- Left e -> throw e- Right (TColumn col) -> case toUnboxedVector @a col of- Left e -> throw e- Right xs -> median' xs---- | Calculates the median of a given column (containing optional values) as a standalone value.-medianMaybe ::- forall a. (Columnable a, Real a) => Expr (Maybe a) -> DataFrame -> Double-medianMaybe (Col name) df =- (median' . optionalToDoubleVector)- (either throw id (columnAsVector (Col @(Maybe a) name) df))-medianMaybe expr df = case interpret @(Maybe a) df expr of- Left e -> throw e- Right (TColumn col) -> case toVector @(Maybe a) col of- Left e -> throw e- Right xs -> (median' . optionalToDoubleVector) xs---- | Calculates the nth percentile of a given column as a standalone value.-percentile ::- forall a.- (Columnable a, Real a, VU.Unbox a) => Int -> Expr a -> DataFrame -> Double-percentile n (Col name) df = case columnAsUnboxedVector (Col @a name) df of- Right xs -> percentile' n xs- Left e -> throw e-percentile n expr df = case interpret df expr of- Left e -> throw e- Right (TColumn col) -> case toUnboxedVector @a col of- Left e -> throw e- Right xs -> percentile' n xs---- | Calculates the nth percentile of a given column as a standalone value.-genericPercentile ::- forall a.- (Columnable a, Ord a) => Int -> Expr a -> DataFrame -> a-genericPercentile n (Col name) df = case columnAsVector (Col @a name) df of- Right xs -> percentileOrd' n xs- Left e -> throw e-genericPercentile n expr df = case interpret df expr of- Left e -> throw e- Right (TColumn col) -> case toVector @a col of- Left e -> throw e- Right xs -> percentileOrd' n xs---- | Calculates the standard deviation of a given column as a standalone value.-standardDeviation ::- forall a. (Columnable a, Real a, VU.Unbox a) => Expr a -> DataFrame -> Double-standardDeviation (Col name) df = case columnAsUnboxedVector (Col @a name) df of- Right xs -> (sqrt . variance') xs- Left e -> throw e-standardDeviation expr df = case interpret df expr of- Left e -> throw e- Right (TColumn col) -> case toUnboxedVector @a col of- Left e -> throw e- Right xs -> (sqrt . variance') xs---- | Calculates the skewness of a given column as a standalone value.-skewness ::- forall a. (Columnable a, Real a, VU.Unbox a) => Expr a -> DataFrame -> Double-skewness (Col name) df = case columnAsUnboxedVector (Col @a name) df of- Right xs -> skewness' xs- Left e -> throw e-skewness expr df = case interpret df expr of- Left e -> throw e- Right (TColumn col) -> case toUnboxedVector @a col of- Left e -> throw e- Right xs -> skewness' xs---- | Calculates the variance of a given column as a standalone value.-variance ::- forall a. (Columnable a, Real a, VU.Unbox a) => Expr a -> DataFrame -> Double-variance (Col name) df = case _getColumnAsDouble name df of- Just xs -> varianceDouble' xs- Nothing -> error "[INTERNAL ERROR] Column is non-numeric"-variance expr df = case interpret df expr of- Left e -> throw e- Right (TColumn col) -> case toUnboxedVector @a col of- Left e -> throw e- Right xs -> variance' xs---- | Calculates the inter-quartile range of a given column as a standalone value.-interQuartileRange ::- forall a. (Columnable a, Real a, VU.Unbox a) => Expr a -> DataFrame -> Double-interQuartileRange (Col name) df = case columnAsUnboxedVector (Col @a name) df of- Right xs -> interQuartileRange' xs- Left e -> throw e-interQuartileRange expr df = case interpret df expr of- Left e -> throw e- Right (TColumn col) -> case toUnboxedVector @a col of- Left e -> throw e- Right xs -> interQuartileRange' xs---- | Calculates the Pearson's correlation coefficient between two given columns as a standalone value.-correlation :: T.Text -> T.Text -> DataFrame -> Maybe Double-correlation first second df = do- f <- _getColumnAsDouble first df- s <- _getColumnAsDouble second df- correlation' f s--_getColumnAsDouble :: T.Text -> DataFrame -> Maybe (VU.Vector Double)-_getColumnAsDouble name df = case getColumn name df of- Just (UnboxedColumn _ (f :: VU.Vector a)) -> case testEquality (typeRep @a) (typeRep @Double) of- Just Refl -> Just f- Nothing -> case sIntegral @a of- STrue -> Just (VU.map fromIntegral f)- SFalse -> case sFloating @a of- STrue -> Just (VU.map realToFrac f)- SFalse -> Nothing- Nothing ->- throw $- ColumnsNotFoundException [name] "_getColumnAsDouble" (M.keys $ columnIndices df)- _ -> Nothing -- Return a type mismatch error here.-{-# INLINE _getColumnAsDouble #-}--optionalToDoubleVector :: (Real a) => V.Vector (Maybe a) -> VU.Vector Double-optionalToDoubleVector =- VU.fromList- . V.foldl'- (\acc e -> if isJust e then realToFrac (fromMaybe 0 e) : acc else acc)- []---- | Calculates the sum of a given column as a standalone value.-sum ::- forall a. (Columnable a, Num a) => Expr a -> DataFrame -> a-sum (Col name) df = case getColumn name df of- Nothing -> throw $ ColumnsNotFoundException [name] "sum" (M.keys $ columnIndices df)- Just ((UnboxedColumn _ (column :: VU.Vector a'))) -> case testEquality (typeRep @a') (typeRep @a) of- Just Refl -> VG.sum column- Nothing -> 0- Just ((BoxedColumn _ (column :: V.Vector a'))) -> case testEquality (typeRep @a') (typeRep @a) of- Just Refl -> VG.sum column- Nothing -> 0-sum expr df = case interpret df expr of- Left e -> throw e- Right (TColumn xs) -> case toVector @a @V.Vector xs of- Left e -> throw e- Right xs' -> VG.sum xs'--{- | /O(n)/ Impute missing values in a column using a derived scalar.--Given--* an expression @f :: 'Expr' b -> 'Expr' b@ that, when interpreted over a- non-nullable column, produces the same value in every row (for example a- mean, median, or other aggregate), and-* a nullable column @'Expr' ('Maybe' b)@--this function:--1. Drops all @Nothing@ values from the target column.-2. Interprets @f@ on the remaining non-null values.-3. Checks that the resulting column contains a single repeated value.-4. Uses that value to impute all @Nothing@s in the original column.--==== __Throws__--* 'DataFrameException' - if the column does not exist, is empty,--==== __Example__-@->>> :set -XOverloadedStrings->>> import qualified DataFrame as D->>> let df =-... D.fromNamedColumns-... [ ("age", D.fromList [Just 10, Nothing, Just 20 :: Maybe Int]) ]->>>->>> -- Impute missing ages with the mean of the observed ages->>> D.imputeWith F.mean "age" df--- age--- ------- 10--- 15--- 20-@--}-instance {-# OVERLAPPING #-} (Columnable b) => ImputeOp (Maybe b) where- runImpute = imputeCore-- runImputeWith f col@(Col columnName) df =- case interpret @b (filterJust columnName df) (f (Col @b columnName)) of- Left e -> throw e- Right (TColumn value) -> case headColumn @b value of- Left e -> throw e- Right h ->- if all (== h) (toList @b value)- then imputeCore col h df- else error "Impute expression returned more than one value"- runImputeWith _ _ df = df--imputeWith ::- forall a.- (ImputeOp a, Columnable (BaseType a)) =>- (Expr (BaseType a) -> Expr (BaseType a)) ->- Expr a ->- DataFrame ->- DataFrame-imputeWith = runImputeWith--applyStatistic ::- (VU.Vector Double -> Double) -> T.Text -> DataFrame -> Maybe Double-applyStatistic f name df = apply =<< _getColumnAsDouble name (filterJust name df)- where- apply col =- let- res = f col- in- if isNaN res then Nothing else pure res-{-# INLINE applyStatistic #-}--applyStatistics ::- (VU.Vector Double -> VU.Vector Double) ->- T.Text ->- DataFrame ->- Maybe (VU.Vector Double)-applyStatistics f name df = fmap f (_getColumnAsDouble name (filterJust name df))---- | Descriptive statistics of the numeric columns.-summarize :: DataFrame -> DataFrame-summarize df =- fold- columnStats- (columnNames df)- ( fromNamedColumns- [- ( "Statistic"- , fromList- [ "Count" :: T.Text- , "Mean"- , "Minimum"- , "25%"- , "Median"- , "75%"- , "Max"- , "StdDev"- , "IQR"- , "Skewness"- ]- )- ]- )- where- columnStats name d =- if all isJust (stats name)- then- insertUnboxedVector- name- (VU.fromList (map (roundTo 2 . fromMaybe 0) $ stats name))- d- else d- stats name =- let- count = fromIntegral . numElements <$> getColumn name df- quantiles = applyStatistics (quantiles' (VU.fromList [0, 1, 2, 3, 4]) 4) name df- min' = flip (VG.!) 0 <$> quantiles- quartile1 = flip (VG.!) 1 <$> quantiles- medianVal = flip (VG.!) 2 <$> quantiles- quartile3 = flip (VG.!) 3 <$> quantiles- max' = flip (VG.!) 4 <$> quantiles- iqr = (-) <$> quartile3 <*> quartile1- doubleColumn col = _getColumnAsDouble col (filterJust col df)- in- [ count- , mean' <$> doubleColumn name- , min'- , quartile1- , medianVal- , quartile3- , max'- , sqrt . variance' <$> doubleColumn name- , iqr- , skewness' <$> doubleColumn name- ]---- | Round a @Double@ to Specified Precision-roundTo :: Int -> Double -> Double-roundTo n x = fromInteger (round $ x * 10 ^ n) / 10.0 ^^ n--toPct2dp :: Double -> String-toPct2dp x- | x < 0.00005 = "<0.01%"- | otherwise = printf "%.2f%%" (x * 100)
− src/DataFrame/Operations/Subset.hs
@@ -1,540 +0,0 @@-{-# LANGUAGE BangPatterns #-}-{-# LANGUAGE CPP #-}-{-# LANGUAGE ConstraintKinds #-}-{-# LANGUAGE ExplicitNamespaces #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE RankNTypes #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}--module DataFrame.Operations.Subset where--import qualified Data.List as L-import qualified Data.Map as M-import qualified Data.Set as S-import qualified Data.Text as T-import qualified Data.Vector as V-import qualified Data.Vector.Generic as VG-import qualified Data.Vector.Unboxed as VU-import qualified Prelude--import Control.Exception (throw)-import Data.Function ((&))-import Data.Maybe (- fromJust,- fromMaybe,- isJust,- isNothing,- )-import Data.Type.Equality (TestEquality (..))-import DataFrame.Errors (- DataFrameException (..),- TypeErrorContext (..),- )-import DataFrame.Internal.Column-import DataFrame.Internal.DataFrame (- DataFrame (..),- derivingExpressions,- empty,- getColumn,- unsafeGetColumn,- )-import DataFrame.Internal.Expression-import DataFrame.Internal.Interpreter-import DataFrame.Operations.Core-import DataFrame.Operations.Merge ()-import DataFrame.Operations.Transformations (apply)-import DataFrame.Operators-import System.Random-import Type.Reflection-import Prelude hiding (filter, take)--#if MIN_VERSION_random(1,3,0)-type SplittableGen g = (SplitGen g, RandomGen g)--splitForStratified :: SplittableGen g => g -> (g, g)-splitForStratified = splitGen-#else-type SplittableGen g = RandomGen g--splitForStratified :: SplittableGen g => g -> (g, g)-splitForStratified = split-#endif---- | O(k * n) Take the first n rows of a DataFrame.-take :: Int -> DataFrame -> DataFrame-take n d = d{columns = V.map (takeColumn n') (columns d), dataframeDimensions = (n', c)}- where- (r, c) = dataframeDimensions d- n' = clip n 0 r---- | O(k * n) Take the last n rows of a DataFrame.-takeLast :: Int -> DataFrame -> DataFrame-takeLast n d =- d- { columns = V.map (takeLastColumn n') (columns d)- , dataframeDimensions = (n', c)- }- where- (r, c) = dataframeDimensions d- n' = clip n 0 r---- | O(k * n) Drop the first n rows of a DataFrame.-drop :: Int -> DataFrame -> DataFrame-drop n d =- d- { columns = V.map (sliceColumn n' (max (r - n') 0)) (columns d)- , dataframeDimensions = (max (r - n') 0, c)- }- where- (r, c) = dataframeDimensions d- n' = clip n 0 r---- | O(k * n) Drop the last n rows of a DataFrame.-dropLast :: Int -> DataFrame -> DataFrame-dropLast n d =- d{columns = V.map (sliceColumn 0 n') (columns d), dataframeDimensions = (n', c)}- where- (r, c) = dataframeDimensions d- n' = clip (r - n) 0 r---- | O(k * n) Take a range of rows of a DataFrame.-range :: (Int, Int) -> DataFrame -> DataFrame-range (start, end) d =- d- { columns = V.map (sliceColumn (clip start 0 r) n') (columns d)- , dataframeDimensions = (n', c)- }- where- (r, c) = dataframeDimensions d- n' = clip (end - start) 0 r--clip :: Int -> Int -> Int -> Int-clip n left right = min right $ max n left--{- | O(n * k) Filter rows by a given condition.--> filter "x" even df--}-filter ::- forall a.- (Columnable a) =>- -- | Column to filter by- Expr a ->- -- | Filter condition- (a -> Bool) ->- -- | Dataframe to filter- DataFrame ->- DataFrame-filter (Col filterColumnName) condition df = case getColumn filterColumnName df of- Nothing ->- throw $- ColumnsNotFoundException [filterColumnName] "filter" (M.keys $ columnIndices df)- Just _col@(BoxedColumn bm (column :: V.Vector b)) ->- -- Check direct type match first, then try Maybe b match for nullable columns- case testEquality (typeRep @a) (typeRep @b) of- Just Refl -> filterByVector filterColumnName column condition df- Nothing -> case (bm, typeRep @a) of- (Just bm', App tMaybe tInner) -> case eqTypeRep tMaybe (typeRep @Maybe) of- Just HRefl -> case testEquality tInner (typeRep @b) of- Just Refl ->- let maybeVec = V.imap (\i v -> if bitmapTestBit bm' i then Just v else Nothing) column- in filterByVector filterColumnName maybeVec condition df- Nothing -> filterByVector filterColumnName column condition df- Nothing -> filterByVector filterColumnName column condition df- _ -> filterByVector filterColumnName column condition df- Just _col@(UnboxedColumn bm (column :: VU.Vector b)) ->- case testEquality (typeRep @a) (typeRep @b) of- Just Refl -> filterByVector filterColumnName column condition df- Nothing -> case (bm, typeRep @a) of- (Just bm', App tMaybe tInner) -> case eqTypeRep tMaybe (typeRep @Maybe) of- Just HRefl -> case testEquality tInner (typeRep @b) of- Just Refl ->- let maybeVec = V.generate (VU.length column) $ \i ->- if bitmapTestBit bm' i then Just (VU.unsafeIndex column i) else Nothing- in filterByVector filterColumnName maybeVec condition df- Nothing -> filterByVector filterColumnName column condition df- Nothing -> filterByVector filterColumnName column condition df- _ -> filterByVector filterColumnName column condition df-filter expr condition df =- let- (TColumn col') = case interpret @a df (normalize expr) of- Left e -> throw e- Right c -> c- indexes = case findIndices condition col' of- Right ixs -> ixs- Left e -> throw e- c' = snd $ dataframeDimensions df- in- df- { columns = V.map (atIndicesStable indexes) (columns df)- , dataframeDimensions = (VU.length indexes, c')- }--filterByVector ::- forall a b v.- (VG.Vector v b, VG.Vector v Int, Columnable a, Columnable b) =>- T.Text -> v b -> (a -> Bool) -> DataFrame -> DataFrame-filterByVector filterColumnName column condition df = case testEquality (typeRep @a) (typeRep @b) of- Nothing ->- throw $- TypeMismatchException- ( MkTypeErrorContext- { userType = Right $ typeRep @a- , expectedType = Right $ typeRep @b- , errorColumnName = Just (T.unpack filterColumnName)- , callingFunctionName = Just "filter"- }- )- Just Refl ->- let- ixs = VG.convert (VG.findIndices condition column)- in- df- { columns = V.map (atIndicesStable ixs) (columns df)- , dataframeDimensions = (VG.length ixs, snd (dataframeDimensions df))- }--{- | O(k) a version of filter where the predicate comes first.--> filterBy even "x" df--}-filterBy :: (Columnable a) => (a -> Bool) -> Expr a -> DataFrame -> DataFrame-filterBy = flip filter--{- | O(k) filters the dataframe with a boolean expression.--> filterWhere (F.col @Int x + F.col y F.> 5) df--}-filterWhere :: Expr Bool -> DataFrame -> DataFrame-filterWhere expr df =- let- (TColumn col') = case interpret @Bool df (normalize expr) of- Left e -> throw e- Right c -> c- indexes = case findIndices id col' of- Right ixs -> ixs- Left e -> throw e- c' = snd $ dataframeDimensions df- in- df- { columns = V.map (atIndicesStable indexes) (columns df)- , dataframeDimensions = (VU.length indexes, c')- }--{- | O(k) removes all rows with `Nothing` in a given column from the dataframe.--> filterJust "col" df--}-filterJust :: T.Text -> DataFrame -> DataFrame-filterJust colName df = case getColumn colName df of- Nothing ->- throw $- ColumnsNotFoundException [colName] "filterJust" (M.keys $ columnIndices df)- Just column | hasMissing column -> case column of- BoxedColumn (Just _) (_col :: V.Vector a) ->- filter (Col @(Maybe a) colName) isJust df & apply @(Maybe a) fromJust colName- UnboxedColumn (Just _) (_col :: VU.Vector a) ->- filter (Col @(Maybe a) colName) isJust df & apply @(Maybe a) fromJust colName- _ -> df- Just _ -> df--{- | O(k) returns all rows with `Nothing` in a give column.--> filterNothing "col" df--}-filterNothing :: T.Text -> DataFrame -> DataFrame-filterNothing colName df = case getColumn colName df of- Nothing ->- throw $- ColumnsNotFoundException [colName] "filterNothing" (M.keys $ columnIndices df)- Just column | hasMissing column -> case column of- BoxedColumn (Just _) (_col :: V.Vector a) -> filter (Col @(Maybe a) colName) isNothing df- UnboxedColumn (Just _) (_col :: VU.Vector a) -> filter (Col @(Maybe a) colName) isNothing df- _ -> df- _ -> df--{- | O(n * k) removes all rows with `Nothing` from the dataframe.--> filterAllJust df--}-filterAllJust :: DataFrame -> DataFrame-filterAllJust df = foldr filterJust df (columnNames df)--{- | O(n * k) keeps any row with a null value.--> filterAllNothing df--}-filterAllNothing :: DataFrame -> DataFrame-filterAllNothing df = foldr filterNothing df (columnNames df)--{- | O(k) cuts the dataframe in a cube of size (a, b) where- a is the length and b is the width.--> cube (10, 5) df--}-cube :: (Int, Int) -> DataFrame -> DataFrame-cube (len, width) = take len . selectBy [ColumnIndexRange (0, width - 1)]--{- | O(n) Selects a number of columns in a given dataframe.--> select ["name", "age"] df--}-select ::- [T.Text] ->- DataFrame ->- DataFrame-select cs df- | L.null cs = empty- | any (`notElem` columnNames df) cs =- throw $- ColumnsNotFoundException- (cs L.\\ columnNames df)- "select"- (columnNames df)- | otherwise =- let result = L.foldl' addKeyValue empty cs- filteredExprs = M.filterWithKey (\k _ -> k `L.elem` cs) (derivingExpressions df)- in result{derivingExpressions = filteredExprs}- where- addKeyValue d k = fromMaybe df $ do- col' <- getColumn k df- pure $ insertColumn k col' d--data SelectionCriteria- = ColumnProperty (Column -> Bool)- | ColumnNameProperty (T.Text -> Bool)- | ColumnTextRange (T.Text, T.Text)- | ColumnIndexRange (Int, Int)- | ColumnName T.Text--{- | Criteria for selecting a column by name.--> selectBy [byName "Age"] df--equivalent to:--> select ["Age"] df--}-byName :: T.Text -> SelectionCriteria-byName = ColumnName--{- | Criteria for selecting columns whose property satisfies given predicate.--> selectBy [byProperty isNumeric] df--}-byProperty :: (Column -> Bool) -> SelectionCriteria-byProperty = ColumnProperty--{- | Criteria for selecting columns whose name satisfies given predicate.--> selectBy [byNameProperty (T.isPrefixOf "weight")] df--}-byNameProperty :: (T.Text -> Bool) -> SelectionCriteria-byNameProperty = ColumnNameProperty--{- | Criteria for selecting columns whose names are in the given lexicographic range (inclusive).--> selectBy [byNameRange ("a", "c")] df--}-byNameRange :: (T.Text, T.Text) -> SelectionCriteria-byNameRange = ColumnTextRange--{- | Criteria for selecting columns whose indices are in the given (inclusive) range.--> selectBy [byIndexRange (0, 5)] df--}-byIndexRange :: (Int, Int) -> SelectionCriteria-byIndexRange = ColumnIndexRange---- | O(n) select columns by column predicate name.-selectBy :: [SelectionCriteria] -> DataFrame -> DataFrame-selectBy xs df = select finalSelection df- where- finalSelection = Prelude.filter (`S.member` columnsWithProperties) (columnNames df)- columnsWithProperties = S.fromList (L.foldl' columnWithProperty [] xs)- columnWithProperty acc (ColumnName colName) = acc ++ [colName]- columnWithProperty acc (ColumnNameProperty f) = acc ++ L.filter f (columnNames df)- columnWithProperty acc (ColumnTextRange (from, to)) =- acc- ++ reverse- (Prelude.dropWhile (to /=) $ reverse $ dropWhile (from /=) (columnNames df))- columnWithProperty acc (ColumnIndexRange (from, to)) = acc ++ Prelude.take (to - from + 1) (Prelude.drop from (columnNames df))- columnWithProperty acc (ColumnProperty f) =- acc- ++ map fst (L.filter (\(_k, v) -> v `elem` ixs) (M.toAscList (columnIndices df)))- where- ixs = V.ifoldl' (\acc' i c -> if f c then i : acc' else acc') [] (columns df)--{- | O(n) inverse of select--> exclude ["Name"] df--}-exclude ::- [T.Text] ->- DataFrame ->- DataFrame-exclude cs df =- let keysToKeep = columnNames df L.\\ cs- in select keysToKeep df--{- | Sample a dataframe. The double parameter must be between 0 and 1 (inclusive).--==== __Example__-@-ghci> import System.Random-ghci> D.sample (mkStdGen 137) 0.1 df--@--}-sample :: (RandomGen g) => g -> Double -> DataFrame -> DataFrame-sample pureGen p df =- let- rand = mkRandom pureGen (fst (dataframeDimensions df)) (0 :: Double) 1- cRand = col @Double "__rand__"- in- df- & insertColumn (name cRand) rand- & filterWhere (cRand .>=. Lit (1 - p))- & exclude [name cRand]--{- | Split a dataset into two. The first in the tuple gets a sample of p (0 <= p <= 1) and the second gets (1 - p). This is useful for creating test and train splits.--==== __Example__-@-ghci> import System.Random-ghci> D.randomSplit (mkStdGen 137) 0.9 df--@--}-randomSplit ::- (RandomGen g) => g -> Double -> DataFrame -> (DataFrame, DataFrame)-randomSplit pureGen p df =- let- rand = mkRandom pureGen (fst (dataframeDimensions df)) (0 :: Double) 1- cRand = col @Double "__rand__"- withRand = df & insertColumn (name cRand) rand- in- ( withRand- & filterWhere (cRand .<=. Lit p)- & exclude [name cRand]- , withRand- & filterWhere- (cRand .>. Lit p)- & exclude [name cRand]- )--{- | Creates n folds of a dataframe.--==== __Example__-@-ghci> import System.Random-ghci> D.kFolds (mkStdGen 137) 5 df--@--}-kFolds :: (RandomGen g) => g -> Int -> DataFrame -> [DataFrame]-kFolds pureGen folds df =- let- rand = mkRandom pureGen (fst (dataframeDimensions df)) (0 :: Double) 1- cRand = col @Double "__rand__"- withRand = df & insertColumn (name cRand) rand- partitionSize = 1 / fromIntegral folds- singleFold n d =- d & filterWhere (cRand .>=. Lit (fromIntegral n * partitionSize))- go (-1) _ = []- go n d =- let- d' = singleFold n d- d'' = d & filterWhere (cRand .<. Lit (fromIntegral n * partitionSize))- in- d' : go (n - 1) d''- in- map (exclude [name cRand]) (go (folds - 1) withRand)---- | Convert any Column to a vector of Text labels (one per row).-columnToTextVec :: Column -> V.Vector T.Text-columnToTextVec (BoxedColumn bm (col' :: V.Vector a)) =- case bm of- Nothing -> case testEquality (typeRep @a) (typeRep @T.Text) of- Just Refl -> col'- Nothing -> V.map (T.pack . show) col'- Just bitmap ->- V.imap (\i x -> if bitmapTestBit bitmap i then T.pack (show x) else "null") col'-columnToTextVec (UnboxedColumn bm col') =- case bm of- Nothing -> V.map (T.pack . show) (V.convert col')- Just bitmap ->- V.generate (VU.length col') $ \i ->- if bitmapTestBit bitmap i then T.pack (show (col' VU.! i)) else "null"---- | Build a map from stringified label to row indices.-groupByIndices :: Column -> M.Map T.Text (VU.Vector Int)-groupByIndices col' =- let textVec = columnToTextVec col'- (grouped, _) =- V.foldl'- (\(!m, !i) key -> (M.insertWith (++) key [i] m, i + 1))- (M.empty, 0)- textVec- in M.map (VU.fromList . L.reverse) grouped---- | Select rows at the given indices from all columns.-rowsAtIndices :: VU.Vector Int -> DataFrame -> DataFrame-rowsAtIndices ixs df =- df- { columns = V.map (atIndicesStable ixs) (columns df)- , dataframeDimensions = (VU.length ixs, snd (dataframeDimensions df))- }--{- | Sample a dataframe, preserving per-stratum proportions.--==== __Example__-@-ghci> import System.Random-ghci> D.stratifiedSample (mkStdGen 42) 0.8 "label" df-@--}-stratifiedSample ::- forall a g.- (SplittableGen g, Columnable a) =>- g -> Double -> Expr a -> DataFrame -> DataFrame-stratifiedSample gen p strataCol df =- let col' = case strataCol of- Col colName -> unsafeGetColumn colName df- _ -> unwrapTypedColumn (either throw id (interpret @a df strataCol))- groups = M.elems (groupByIndices col')- go _ [] = mempty- go g (ixs : rest) =- let stratum = rowsAtIndices ixs df- (g1, g2) = splitForStratified g- in sample g1 p stratum <> go g2 rest- in go gen groups--{- | Split a dataframe into two, preserving per-stratum proportions.--==== __Example__-@-ghci> import System.Random-ghci> D.stratifiedSplit (mkStdGen 42) 0.8 "label" df-@--}-stratifiedSplit ::- forall a g.- (SplittableGen g, Columnable a) =>- g -> Double -> Expr a -> DataFrame -> (DataFrame, DataFrame)-stratifiedSplit gen p strataCol df =- let col' = case strataCol of- Col colName -> unsafeGetColumn colName df- _ -> unwrapTypedColumn (either throw id (interpret @a df strataCol))- groups = M.elems (groupByIndices col')- go _ [] = (mempty, mempty)- go g (ixs : rest) =- let stratum = rowsAtIndices ixs df- (g1, g2) = splitForStratified g- (tr, va) = randomSplit g1 p stratum- (trAcc, vaAcc) = go g2 rest- in (tr <> trAcc, va <> vaAcc)- in go gen groups
− src/DataFrame/Operations/Transformations.hs
@@ -1,244 +0,0 @@-{-# LANGUAGE ConstrainedClassMethods #-}-{-# LANGUAGE ExplicitNamespaces #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE FlexibleInstances #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE RankNTypes #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}-{-# LANGUAGE UndecidableInstances #-}-{-# LANGUAGE UndecidableSuperClasses #-}--module DataFrame.Operations.Transformations where--import qualified Data.List as L-import qualified Data.Map as M-import qualified Data.Text as T-import qualified Data.Vector as V--import Control.Exception (throw)-import Data.Maybe-import DataFrame.Errors (DataFrameException (..), TypeErrorContext (..))-import DataFrame.Internal.Column (- Columnable,- TypedColumn (..),- hasMissing,- ifoldrColumn,- imapColumn,- mapColumn,- )-import DataFrame.Internal.DataFrame (DataFrame (..), getColumn)-import DataFrame.Internal.Expression-import DataFrame.Internal.Interpreter-import DataFrame.Internal.Nullable (BaseType)-import DataFrame.Operations.Core---- | O(k) Apply a function to a given column in a dataframe.-apply ::- forall b c.- (Columnable b, Columnable c) =>- -- | function to apply- (b -> c) ->- -- | Column name- T.Text ->- -- | DataFrame to apply operation to- DataFrame ->- DataFrame-apply f columnName d = case safeApply f columnName d of- Left (TypeMismatchException context) ->- throw $ TypeMismatchException (context{callingFunctionName = Just "apply"})- Left exception -> throw exception- Right df -> df---- | O(k) Safe version of the apply function. Returns (instead of throwing) the error.-safeApply ::- forall b c.- (Columnable b, Columnable c) =>- -- | function to apply- (b -> c) ->- -- | Column name- T.Text ->- -- | DataFrame to apply operation to- DataFrame ->- Either DataFrameException DataFrame-safeApply f columnName d = case getColumn columnName d of- Nothing ->- Left $ ColumnsNotFoundException [columnName] "apply" (M.keys $ columnIndices d)- Just column -> do- column' <- mapColumn f column- pure $ insertColumn columnName column' d--{- | O(k) Apply a function to an expression in a dataframe and-add the result into `alias` column.--}-derive :: forall a. (Columnable a) => T.Text -> Expr a -> DataFrame -> DataFrame-derive name expr df = case interpret @a df (normalize expr) of- Left e -> throw e- Right (TColumn value) ->- (insertColumn name value df)- { derivingExpressions = M.insert name (UExpr expr) (derivingExpressions df)- }--{- | O(k) Apply a function to an expression in a dataframe and-add the result into `alias` column but--==== __Examples__-->>> (z, df') = deriveWithExpr "z" (F.col @Int "x" + F.col "y") df->>> filterWhere (z .>= 50)--}-deriveWithExpr ::- forall a. (Columnable a) => T.Text -> Expr a -> DataFrame -> (Expr a, DataFrame)-deriveWithExpr name expr df = case interpret @a df (normalize expr) of- Left e -> throw e- Right (TColumn value) ->- ( Col name- , (insertColumn name value df)- { derivingExpressions = M.insert name (UExpr expr) (derivingExpressions df)- }- )--deriveMany :: [NamedExpr] -> DataFrame -> DataFrame-deriveMany exprs df =- let- f (name, UExpr (expr :: Expr a)) d =- case interpret @a df expr of- Left e -> throw e- Right (TColumn value) -> insertColumn name value d- in- fold f exprs df---- | O(k * n) Apply a function to given column names in a dataframe.-applyMany ::- (Columnable b, Columnable c) =>- (b -> c) ->- [T.Text] ->- DataFrame ->- DataFrame-applyMany f names df = L.foldl' (flip (apply f)) df names---- | O(k) Convenience function that applies to an int column.-applyInt ::- (Columnable b) =>- -- | function to apply- (Int -> b) ->- -- | Column name- T.Text ->- -- | DataFrame to apply operation to- DataFrame ->- DataFrame-applyInt = apply---- | O(k) Convenience function that applies to an double column.-applyDouble ::- (Columnable b) =>- -- | function to apply- (Double -> b) ->- -- | Column name- T.Text ->- -- | DataFrame to apply operation to- DataFrame ->- DataFrame-applyDouble = apply--{- | O(k * n) Apply a function to a column only if there is another column-value that matches the given criterion.--> applyWhere (<20) "Age" (const "Gen-Z") "Generation" df--}-applyWhere ::- forall a b.- (Columnable a, Columnable b) =>- -- | Filter condition- (a -> Bool) ->- -- | Criterion Column- T.Text ->- -- | function to apply- (b -> b) ->- -- | Column name- T.Text ->- -- | DataFrame to apply operation to- DataFrame ->- DataFrame-applyWhere condition filterColumnName f columnName df = case getColumn filterColumnName df of- Nothing ->- throw $- ColumnsNotFoundException- [filterColumnName]- "applyWhere"- (M.keys $ columnIndices df)- Just column -> case ifoldrColumn- (\i val acc -> if condition val then V.cons i acc else acc)- V.empty- column of- Left e -> throw e- Right indexes ->- if V.null indexes- then df- else L.foldl' (\d i -> applyAtIndex i f columnName d) df indexes---- | O(k) Apply a function to the column at a given index.-applyAtIndex ::- forall a.- (Columnable a) =>- -- | Index- Int ->- -- | function to apply- (a -> a) ->- -- | Column name- T.Text ->- -- | DataFrame to apply operation to- DataFrame ->- DataFrame-applyAtIndex i f columnName df = case getColumn columnName df of- Nothing ->- throw $- ColumnsNotFoundException [columnName] "applyAtIndex" (M.keys $ columnIndices df)- Just column -> case imapColumn (\index value -> if index == i then f value else value) column of- Left e -> throw e- Right column' -> insertColumn columnName column' df---- | Core impute implementation for nullable columns. Silently no-ops on non-nullable columns.-imputeCore ::- forall b.- (Columnable b) =>- Expr (Maybe b) ->- b ->- DataFrame ->- DataFrame-imputeCore (Col columnName) value df = case getColumn columnName df of- Nothing ->- throw $- ColumnsNotFoundException [columnName] "impute" (M.keys $ columnIndices df)- Just col | hasMissing col -> case safeApply (fromMaybe value) columnName df of- Left (TypeMismatchException context) -> throw $ TypeMismatchException (context{callingFunctionName = Just "impute"})- Left exception -> throw exception- Right res -> res- _ -> df-imputeCore _ _ df = df--class (Columnable a) => ImputeOp a where- runImpute :: Expr a -> BaseType a -> DataFrame -> DataFrame- runImputeWith ::- (Columnable (BaseType a)) =>- (Expr (BaseType a) -> Expr (BaseType a)) ->- Expr a ->- DataFrame ->- DataFrame--instance {-# OVERLAPPABLE #-} (Columnable a) => ImputeOp a where- runImpute _ _ df = df- runImputeWith _ _ df = df--{- | Replace all instances of `Nothing` in a column with the given value.-When the column is already non-nullable, this is a silent no-op.--}-impute ::- forall a.- (ImputeOp a) =>- Expr a ->- BaseType a ->- DataFrame ->- DataFrame-impute = runImpute
− src/DataFrame/Operations/Typing.hs
@@ -1,470 +0,0 @@-{-# LANGUAGE AllowAmbiguousTypes #-}-{-# LANGUAGE BangPatterns #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}--module DataFrame.Operations.Typing where--import qualified Data.Map as M-import qualified Data.Text as T-import qualified Data.Vector as V-import qualified Data.Vector.Mutable as VM-import qualified Data.Vector.Unboxed as VU-import qualified Data.Vector.Unboxed.Mutable as VUM--import Control.Applicative (asum)-import Control.Monad (join)-import Control.Monad.ST (runST)-import Data.Maybe (fromMaybe)-import qualified Data.Proxy as P-import Data.Time-import Data.Type.Equality (TestEquality (..))-import DataFrame.Internal.Column (- Bitmap,- Column (..),- Columnable,- bitmapTestBit,- ensureOptional,- finalizeParseResult,- fromVector,- )-import DataFrame.Internal.DataFrame (DataFrame (..), unsafeGetColumn)-import DataFrame.Internal.Parsing-import DataFrame.Internal.Schema-import DataFrame.Operations.Core-import Text.Read-import Type.Reflection--type DateFormat = String--{- | How parse failures are surfaced in the resulting column.--* 'NoSafeRead' — strict parsing: failures throw (via 'read').-* 'MaybeRead' — failures become 'Nothing'; columns are wrapped as @Maybe a@.-* 'EitherRead' — failures become @Left rawText@; columns are wrapped as- @Either Text a@, preserving the original input so callers can inspect it.--}-data SafeReadMode- = NoSafeRead- | MaybeRead- | EitherRead- deriving (Eq, Show, Read)---- | Options controlling how text columns are parsed into typed values.-data ParseOptions = ParseOptions- { missingValues :: [T.Text]- -- ^ Values to treat as @Nothing@ when the effective mode is 'MaybeRead'.- , sampleSize :: Int- -- ^ Number of rows to inspect when inferring a column's type (0 = all rows).- , parseSafe :: SafeReadMode- {- ^ Default 'SafeReadMode' applied to every column that does not have an- entry in 'parseSafeOverrides'. 'NoSafeRead' only treats empty strings as- missing; 'MaybeRead' additionally treats 'missingValues' and nullish- strings as @Nothing@; 'EitherRead' wraps the resulting column as- @Either Text a@ with the raw input preserved on failure.- -}- , parseSafeOverrides :: [(T.Text, SafeReadMode)]- {- ^ Per-column overrides. When a column name is present here, its value- takes precedence over 'parseSafe'. Typical use: strict IDs- (@NoSafeRead@) alongside lenient fields (@MaybeRead@/@EitherRead@).- -}- , parseDateFormat :: DateFormat- -- ^ Date format string as accepted by "Data.Time.Format" (e.g. @\"%Y-%m-%d\"@).- }--{- | Sensible out-of-the-box parse options: infer from the first 100 rows,- treat common nullish strings as missing, and expect ISO 8601 dates.--}-defaultParseOptions :: ParseOptions-defaultParseOptions =- ParseOptions- { missingValues = []- , sampleSize = 100- , parseSafe = MaybeRead- , parseSafeOverrides = []- , parseDateFormat = "%Y-%m-%d"- }--{- | Resolve a column's effective 'SafeReadMode': the override if present,-otherwise the default.--}-effectiveSafeRead ::- SafeReadMode -> [(T.Text, SafeReadMode)] -> T.Text -> SafeReadMode-effectiveSafeRead def overrides name = fromMaybe def (lookup name overrides)--parseDefaults :: ParseOptions -> DataFrame -> DataFrame-parseDefaults opts df = df{columns = V.imap forCol (columns df)}- where- -- Index -> column name: reverse the columnIndices map once.- nameAt =- let inverted = M.fromList [(i, n) | (n, i) <- M.toList (columnIndices df)]- in \i -> M.findWithDefault "" i inverted- forCol i col =- let mode =- effectiveSafeRead- (parseSafe opts)- (parseSafeOverrides opts)- (nameAt i)- in parseDefault opts{parseSafe = mode, parseSafeOverrides = []} col--parseDefault :: ParseOptions -> Column -> Column-parseDefault opts (BoxedColumn Nothing (c :: V.Vector a)) =- case (typeRep @a) `testEquality` (typeRep @T.Text) of- Nothing -> case (typeRep @a) `testEquality` (typeRep @String) of- Just Refl -> parseFromExamples opts (V.map T.pack c)- Nothing -> BoxedColumn Nothing c- Just Refl -> parseFromExamples opts c-parseDefault opts (BoxedColumn (Just bm) (c :: V.Vector a)) =- case (typeRep @a) `testEquality` (typeRep @T.Text) of- Nothing -> case (typeRep @a) `testEquality` (typeRep @String) of- Just Refl ->- parseFromExamples- opts- (V.imap (\i x -> if bitmapTestBit bm i then T.pack x else "") c)- Nothing -> BoxedColumn (Just bm) c- Just Refl ->- parseFromExamples opts (V.imap (\i x -> if bitmapTestBit bm i then x else "") c)-parseDefault _ column = column--parseFromExamples :: ParseOptions -> V.Vector T.Text -> Column-parseFromExamples opts cols =- let isNull = case parseSafe opts of- NoSafeRead -> T.null- _ -> isNullishOrMissing (missingValues opts)- -- `examples` is small (≤ sampleSize, default 100), so the- -- Maybe-wrap allocation here is ignorable. The full-column- -- equivalent (`asMaybeText = V.map ... cols`) has been removed:- -- handlers now walk `cols` directly with `isNull`.- examples = V.map (classify isNull) (V.take (sampleSize opts) cols)- dfmt = parseDateFormat opts- assumption = makeParsingAssumption dfmt examples- in case parseSafe opts of- EitherRead -> handleEitherAssumption dfmt assumption cols- mode ->- let result = case assumption of- BoolAssumption -> handleBoolAssumption isNull cols- IntAssumption -> handleIntAssumption isNull cols- DoubleAssumption -> handleDoubleAssumption isNull cols- TextAssumption -> handleTextAssumption isNull cols- DateAssumption -> handleDateAssumption dfmt isNull cols- NoAssumption -> handleNoAssumption dfmt isNull cols- in if mode == MaybeRead then ensureOptional result else result- where- classify p t = if p t then Nothing else Just t--{- | For 'EitherRead' mode: take the chosen parsing assumption and produce an-@Either Text a@ column. Successful parses become @Right@; any row that fails-to parse as the chosen type (including null/missing cells) becomes @Left@-carrying the raw input text verbatim.--}-handleEitherAssumption ::- DateFormat -> ParsingAssumption -> V.Vector T.Text -> Column-handleEitherAssumption dfmt assumption raw = case assumption of- BoolAssumption -> fromVector (V.map (toEither readBool) raw)- IntAssumption -> fromVector (V.map (toEither readInt) raw)- DoubleAssumption -> fromVector (V.map (toEither readDouble) raw)- DateAssumption -> fromVector (V.map (toEither (parseTimeOpt dfmt)) raw)- -- TextAssumption and NoAssumption degenerate to Either Text Text; treat- -- empty strings as Left "" so the convention (Left = missing/failure) stays- -- consistent across column types.- TextAssumption -> fromVector (V.map textToEither raw)- NoAssumption -> fromVector (V.map textToEither raw)- where- toEither :: (T.Text -> Maybe a) -> T.Text -> Either T.Text a- toEither p t = maybe (Left t) Right (p t)-- textToEither :: T.Text -> Either T.Text T.Text- textToEither t = if T.null t then Left t else Right t--parseUnboxedColumnWithPred ::- forall src a.- (VU.Unbox a) =>- a ->- (src -> Bool) ->- (src -> Maybe a) ->- V.Vector src ->- Maybe (Maybe Bitmap, VU.Vector a)-parseUnboxedColumnWithPred nullValue isNull parser vec = runST $ do- let n = V.length vec- values <- VUM.unsafeNew n- vmask <- VUM.unsafeNew n- let go !i !anyNull- | i >= n = finalizeParseResult values vmask anyNull- | otherwise =- let !src = V.unsafeIndex vec i- in if isNull src- then do- VUM.unsafeWrite vmask i 0- VUM.unsafeWrite values i nullValue- go (i + 1) True- else case parser src of- Just v -> do- VUM.unsafeWrite vmask i 1- VUM.unsafeWrite values i v- go (i + 1) anyNull- Nothing -> return Nothing- go 0 False-{-# INLINE parseUnboxedColumnWithPred #-}---- | Wrap a successful 'parseUnboxedColumnWithPred' result as a 'Column'.-unboxedOrFallback ::- (Columnable a, VU.Unbox a) =>- Maybe (Maybe Bitmap, VU.Vector a) ->- Column ->- Column-unboxedOrFallback (Just (mbm, vec)) _ = UnboxedColumn mbm vec-unboxedOrFallback Nothing fallback = fallback--handleBoolAssumption :: (T.Text -> Bool) -> V.Vector T.Text -> Column-handleBoolAssumption isNull cols =- unboxedOrFallback- (parseUnboxedColumnWithPred False isNull readBool cols)- (handleTextAssumption isNull cols)--handleIntAssumption :: (T.Text -> Bool) -> V.Vector T.Text -> Column-handleIntAssumption isNull cols =- case parseUnboxedColumnWithPred 0 isNull readInt cols of- Just (mbm, vec) -> UnboxedColumn mbm vec- Nothing ->- unboxedOrFallback- (parseUnboxedColumnWithPred 0 isNull readDouble cols)- (handleTextAssumption isNull cols)--handleDoubleAssumption :: (T.Text -> Bool) -> V.Vector T.Text -> Column-handleDoubleAssumption isNull cols =- unboxedOrFallback- (parseUnboxedColumnWithPred 0 isNull readDouble cols)- (handleTextAssumption isNull cols)--{- | Text columns: no parse, just null-marking. When the whole column-is non-null we return a plain 'V.Vector T.Text'; otherwise we emit a-@V.Vector (Maybe T.Text)@ the same shape the old code produced.--}-handleTextAssumption :: (T.Text -> Bool) -> V.Vector T.Text -> Column-handleTextAssumption isNull cols- | V.any isNull cols =- fromVector- (V.map (\t -> if isNull t then Nothing else Just t) cols)- | otherwise = fromVector cols--{- | Date: single parse pass, boxed because 'Day' is not unboxable.-Bails to 'handleTextAssumption' the moment a non-null cell fails to-parse as a 'Day'. Still avoids the outer @V.Vector (Maybe T.Text)@-allocation — we walk @cols@ directly with @isNull@.--}-handleDateAssumption ::- DateFormat -> (T.Text -> Bool) -> V.Vector T.Text -> Column-handleDateAssumption dateFormat isNull cols =- case parseBoxedMaybeColumn isNull (parseTimeOpt dateFormat) cols of- Just (anyNull, vec)- -- `vec :: V.Vector (Maybe Day)`. If no nulls, strip the- -- outer 'Maybe' (every cell is guaranteed 'Just') so the- -- column type stays 'Day' rather than becoming 'Maybe Day'.- | anyNull -> fromVector vec- | otherwise -> fromVector (V.mapMaybe id vec)- Nothing -> handleTextAssumption isNull cols--parseBoxedMaybeColumn ::- (T.Text -> Bool) ->- (T.Text -> Maybe a) ->- V.Vector T.Text ->- Maybe (Bool, V.Vector (Maybe a))-parseBoxedMaybeColumn isNull parser cols = runST $ do- let n = V.length cols- out <- VM.new n- let loop !i !anyNull- | i >= n = do- frozen <- V.unsafeFreeze out- return (Just (anyNull, frozen))- | otherwise =- let !t = V.unsafeIndex cols i- in if isNull t- then do- VM.unsafeWrite out i Nothing- loop (i + 1) True- else case parser t of- Just v -> do- VM.unsafeWrite out i (Just v)- loop (i + 1) anyNull- Nothing -> return Nothing- loop 0 False--handleNoAssumption ::- DateFormat -> (T.Text -> Bool) -> V.Vector T.Text -> Column-handleNoAssumption dateFormat isNull cols- -- Only reached when the 100-row sample was all-null. Try each- -- concrete type in turn; fall back to Text otherwise.- | V.all isNull cols =- fromVector (V.map (const (Nothing :: Maybe T.Text)) cols)- | Just (mbm, vec) <- parseUnboxedColumnWithPred False isNull readBool cols =- UnboxedColumn mbm vec- | Just (mbm, vec) <- parseUnboxedColumnWithPred 0 isNull readInt cols =- UnboxedColumn mbm vec- | Just (mbm, vec) <- parseUnboxedColumnWithPred 0 isNull readDouble cols =- UnboxedColumn mbm vec- | otherwise = case parseBoxedMaybeColumn isNull (parseTimeOpt dateFormat) cols of- Just (anyNull, vec)- -- `vec :: V.Vector (Maybe Day)`. If no nulls, strip the- -- outer 'Maybe' (every cell is guaranteed 'Just') so the- -- column type stays 'Day' rather than becoming 'Maybe Day'.- | anyNull -> fromVector vec- | otherwise -> fromVector (V.mapMaybe id vec)- Nothing -> handleTextAssumption isNull cols--{- | Predicate matching what 'parseSafe == NoSafeRead' previously used:-only empty strings are treated as missing.--We still expose 'convertNullish' \/ 'convertOnlyEmpty' below because-other parts of the library reference them, but neither is used by-'parseFromExamples' any longer.--}-isNullishOrMissing :: [T.Text] -> T.Text -> Bool-isNullishOrMissing missing v = isNullish v || v `elem` missing--convertNullish :: [T.Text] -> T.Text -> Maybe T.Text-convertNullish missing v = if isNullish v || v `elem` missing then Nothing else Just v--convertOnlyEmpty :: T.Text -> Maybe T.Text-convertOnlyEmpty v = if v == "" then Nothing else Just v--parseTimeOpt :: DateFormat -> T.Text -> Maybe Day-parseTimeOpt dateFormat s =- parseTimeM {- Accept leading/trailing whitespace -}- True- defaultTimeLocale- dateFormat- (T.unpack s)--unsafeParseTime :: DateFormat -> T.Text -> Day-unsafeParseTime dateFormat s =- parseTimeOrError {- Accept leading/trailing whitespace -}- True- defaultTimeLocale- dateFormat- (T.unpack s)--hasNullValues :: (Eq a) => V.Vector (Maybe a) -> Bool-hasNullValues = V.any (== Nothing)--vecSameConstructor :: V.Vector (Maybe a) -> V.Vector (Maybe b) -> Bool-vecSameConstructor xs ys = (V.length xs == V.length ys) && V.and (V.zipWith hasSameConstructor xs ys)- where- hasSameConstructor :: Maybe a -> Maybe b -> Bool- hasSameConstructor (Just _) (Just _) = True- hasSameConstructor Nothing Nothing = True- hasSameConstructor _ _ = False--makeParsingAssumption ::- DateFormat -> V.Vector (Maybe T.Text) -> ParsingAssumption-makeParsingAssumption dateFormat asMaybeText- -- All the examples are "NA", "Null", "", so we can't make any shortcut- -- assumptions and just have to go the long way.- | V.all (== Nothing) asMaybeText = NoAssumption- -- After accounting for nulls, parsing for Ints and Doubles results in the- -- same corresponding positions of Justs and Nothings, so we assume- -- that the best way to parse is Int- | vecSameConstructor asMaybeText asMaybeBool = BoolAssumption- | vecSameConstructor asMaybeText asMaybeInt- && vecSameConstructor asMaybeText asMaybeDouble =- IntAssumption- -- After accounting for nulls, the previous condition fails, so some (or none) can be parsed as Ints- -- and some can be parsed as Doubles, so we make the assumpotion of doubles.- | vecSameConstructor asMaybeText asMaybeDouble = DoubleAssumption- -- After accounting for nulls, parsing for Dates results in the same corresponding- -- positions of Justs and Nothings, so we assume that the best way to parse is Date.- | vecSameConstructor asMaybeText asMaybeDate = DateAssumption- | otherwise = TextAssumption- where- asMaybeBool = V.map (>>= readBool) asMaybeText- asMaybeInt = V.map (>>= readInt) asMaybeText- asMaybeDouble = V.map (>>= readDouble) asMaybeText- asMaybeDate = V.map (>>= parseTimeOpt dateFormat) asMaybeText--data ParsingAssumption- = BoolAssumption- | IntAssumption- | DoubleAssumption- | DateAssumption- | NoAssumption- | TextAssumption--{- | Re-type columns of a 'DataFrame' according to the supplied schema map.-The caller provides a @resolveMode@ function that maps a column name to its-'SafeReadMode' — typically built from a global default plus an overrides map-via 'effectiveSafeRead'.--}-parseWithTypes ::- (T.Text -> SafeReadMode) ->- M.Map T.Text SchemaType ->- DataFrame ->- DataFrame-parseWithTypes resolveMode ts df- | M.null ts = df- | otherwise =- M.foldrWithKey- (\k v d -> insertColumn k (asType (resolveMode k) v (unsafeGetColumn k d)) d)- df- ts- where- -- \| Re-parse a plain (non-Maybe, non-Either) target type according to the- -- 'SafeReadMode'. @toStr@ converts column elements to a 'String' ready for- -- 'Read'.- plainType ::- forall a b.- (Columnable a, Read a) =>- SafeReadMode -> V.Vector b -> (b -> String) -> Column- plainType mode col toStr = case mode of- NoSafeRead -> fromVector (V.map ((read @a) . toStr) col)- MaybeRead -> fromVector (V.map ((readMaybe @a) . toStr) col)- EitherRead -> fromVector (V.map ((readEitherRaw @a) . toStr) col)-- asType :: SafeReadMode -> SchemaType -> Column -> Column- asType mode (SType (_ :: P.Proxy a)) c@(BoxedColumn _ (col :: V.Vector b)) = case typeRep @a of- App t1 _t2 -> case eqTypeRep t1 (typeRep @Maybe) of- Just HRefl -> case testEquality (typeRep @a) (typeRep @b) of- Just Refl -> c- Nothing -> case testEquality (typeRep @T.Text) (typeRep @b) of- Just Refl -> fromVector (V.map (join . (readAsMaybe @a) . T.unpack) col)- Nothing -> fromVector (V.map (join . (readAsMaybe @a) . show) col)- Nothing -> case t1 of- App t1' _t2' -> case eqTypeRep t1' (typeRep @Either) of- Just HRefl -> case testEquality (typeRep @a) (typeRep @b) of- Just Refl -> c- Nothing -> case testEquality (typeRep @T.Text) (typeRep @b) of- Just Refl -> fromVector (V.map ((readAsEither @a) . T.unpack) col)- Nothing -> fromVector (V.map ((readAsEither @a) . show) col)- Nothing -> case testEquality (typeRep @a) (typeRep @b) of- Just Refl -> c- Nothing -> case testEquality (typeRep @T.Text) (typeRep @b) of- Just Refl -> plainType @a mode col T.unpack- Nothing -> plainType @a mode col show- _ -> c- _ -> case testEquality (typeRep @a) (typeRep @b) of- Just Refl -> c- Nothing -> case testEquality (typeRep @T.Text) (typeRep @b) of- Just Refl -> plainType @a mode col T.unpack- Nothing -> plainType @a mode col show- asType _ _ c = c--readAsMaybe :: (Read a) => String -> Maybe a-readAsMaybe s- | null s = Nothing- | otherwise = readMaybe $ "Just " <> s--readAsEither :: (Read a) => String -> a-readAsEither v = case asum [readMaybe $ "Left " <> s, readMaybe $ "Right " <> s] of- Nothing -> error $ "Couldn't read value: " <> s- Just v' -> v'- where- s = if null v then "\"\"" else v--{- | Try 'readMaybe'; on failure return @Left raw@ where @raw@ is the original-input text. Used by 'parseWithTypes' under 'EitherRead'.--}-readEitherRaw :: forall a. (Read a) => String -> Either T.Text a-readEitherRaw s = case readMaybe s of- Just v -> Right v- Nothing -> Left (T.pack s)
− src/DataFrame/Operators.hs
@@ -1,329 +0,0 @@-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE TypeFamilies #-}-{-# LANGUAGE TypeOperators #-}--module DataFrame.Operators where--import Data.Function ((&))-import qualified Data.Text as T-import DataFrame.Internal.Column (Columnable)-import DataFrame.Internal.Expression (- BinaryOp (- MkBinaryOp,- binaryCommutative,- binaryFn,- binaryName,- binaryPrecedence,- binarySymbol- ),- Expr (Binary, Col, If, Lit, Unary),- NamedExpr,- UExpr (UExpr),- UnaryOp (MkUnaryOp, unaryFn, unaryName, unarySymbol),- )-import DataFrame.Internal.Nullable (- BaseType,- DivWidenOp,- NullCmpResult,- NullLift2Op (applyNull2),- NullableCmpOp (nullCmpOp),- NumericWidenOp,- WidenResult,- WidenResultDiv,- divArithOp,- widenArithOp,- widenCmpOp,- )-import DataFrame.Internal.Types (Promote, PromoteDiv)--infixr 8 .^^, .^^., .^, .^.-infixl 7 .*, ./, .*., ./.-infixl 6 .+, .-, .+., .-.-infix 4 .==, .==., .<, .<., .<=, .<=., .>=, .>=., .>, .>., ./=, ./=.-infixr 3 .&&, .&&.-infixr 2 .||, .||.-infixr 0 .=--(|>) :: a -> (a -> b) -> b-(|>) = (&)--as :: (Columnable a) => Expr a -> T.Text -> NamedExpr-as expr colName = (colName, UExpr expr)--name :: (Show a) => Expr a -> T.Text-name (Col n) = n-name other =- error $- "You must call `name` on a column reference. Not the expression: " ++ show other--col :: (Columnable a) => T.Text -> Expr a-col = Col--ifThenElse :: (Columnable a) => Expr Bool -> Expr a -> Expr a -> Expr a-ifThenElse = If--lit :: (Columnable a) => a -> Expr a-lit = Lit--(.=) :: (Columnable a) => T.Text -> Expr a -> NamedExpr-(.=) = flip as--liftDecorated ::- (Columnable a, Columnable b) =>- (a -> b) -> T.Text -> Maybe T.Text -> Expr a -> Expr b-liftDecorated f opName rep = Unary (MkUnaryOp{unaryFn = f, unaryName = opName, unarySymbol = rep})--lift2Decorated ::- (Columnable c, Columnable b, Columnable a) =>- (c -> b -> a) ->- T.Text ->- Maybe T.Text ->- Bool ->- Int ->- Expr c ->- Expr b ->- Expr a-lift2Decorated f opName rep comm prec =- Binary- ( MkBinaryOp- { binaryFn = f- , binaryName = opName- , binarySymbol = rep- , binaryCommutative = comm- , binaryPrecedence = prec- }- )--(.==.) ::- (Columnable a, Eq a) =>- Expr a ->- Expr a ->- Expr Bool-(.==.) = lift2Decorated (==) "eq" (Just ".==.") True 4--(./=.) ::- (Columnable a, Eq a) =>- Expr a ->- Expr a ->- Expr Bool-(./=.) = lift2Decorated (/=) "neq" (Just "./=.") True 4--(.<.) ::- (Columnable a, Ord a) =>- Expr a ->- Expr a ->- Expr Bool-(.<.) = lift2Decorated (<) "lt" (Just ".<.") False 4--(.>.) ::- (Columnable a, Ord a) =>- Expr a ->- Expr a ->- Expr Bool-(.>.) = lift2Decorated (>) "gt" (Just ".>.") False 4--(.<=.) ::- (Columnable a, Ord a) =>- Expr a ->- Expr a ->- Expr Bool-(.<=.) = lift2Decorated (<=) "leq" (Just ".<=.") False 4--(.>=.) ::- (Columnable a, Ord a) =>- Expr a ->- Expr a ->- Expr Bool-(.>=.) = lift2Decorated (>=) "geq" (Just ".>=.") False 4--(.+.) :: (Columnable a, Num a) => Expr a -> Expr a -> Expr a-(.+.) = (+)--(.-.) :: (Columnable a, Num a) => Expr a -> Expr a -> Expr a-(.-.) = (-)--(.*.) :: (Columnable a, Num a) => Expr a -> Expr a -> Expr a-(.*.) = (*)--(./.) :: (Columnable a, Fractional a) => Expr a -> Expr a -> Expr a-(./.) = (/)---- Nullable-aware arithmetic operators--{- | Nullable-aware addition. Works for all combinations of nullable\/non-nullable operands.-@col \@Int "x" .+ col \@(Maybe Int) "y" -- :: Expr (Maybe Int)@--}-(.+) ::- ( NumericWidenOp (BaseType a) (BaseType b)- , NullLift2Op a b (Promote (BaseType a) (BaseType b)) (WidenResult a b)- , Num (Promote (BaseType a) (BaseType b))- ) =>- Expr a ->- Expr b ->- Expr (WidenResult a b)-(.+) = lift2Decorated (applyNull2 (widenArithOp (+))) "nulladd" (Just ".+") True 6---- | Nullable-aware subtraction.-(.-) ::- ( NumericWidenOp (BaseType a) (BaseType b)- , NullLift2Op a b (Promote (BaseType a) (BaseType b)) (WidenResult a b)- , Num (Promote (BaseType a) (BaseType b))- ) =>- Expr a ->- Expr b ->- Expr (WidenResult a b)-(.-) = lift2Decorated (applyNull2 (widenArithOp (-))) "nullsub" (Just ".-") False 6---- | Nullable-aware multiplication.-(.*) ::- ( NumericWidenOp (BaseType a) (BaseType b)- , NullLift2Op a b (Promote (BaseType a) (BaseType b)) (WidenResult a b)- , Num (Promote (BaseType a) (BaseType b))- ) =>- Expr a ->- Expr b ->- Expr (WidenResult a b)-(.*) = lift2Decorated (applyNull2 (widenArithOp (*))) "nullmul" (Just ".*") True 7---- | Nullable-aware division. Integral operands are promoted to Double.-(./) ::- ( DivWidenOp (BaseType a) (BaseType b)- , NullLift2Op a b (PromoteDiv (BaseType a) (BaseType b)) (WidenResultDiv a b)- , Fractional (PromoteDiv (BaseType a) (BaseType b))- ) =>- Expr a ->- Expr b ->- Expr (WidenResultDiv a b)-(./) = lift2Decorated (applyNull2 (divArithOp (/))) "nulldiv" (Just "./") False 7---- Nullable-aware comparison operators (three-valued logic: Nothing if either operand is Nothing)--{- | Nullable-aware equality. Widens numeric operands to their common type,-so @Expr Double .== Expr Int@ typechecks. Returns @Maybe Bool@ when either-operand is nullable.--}-(.==) ::- ( NumericWidenOp (BaseType a) (BaseType b)- , NullLift2Op a b Bool (NullCmpResult a b)- , Eq (Promote (BaseType a) (BaseType b))- ) =>- Expr a ->- Expr b ->- Expr (NullCmpResult a b)-(.==) = lift2Decorated (applyNull2 (widenCmpOp (==))) "eq" (Just ".==") True 4---- | Nullable-aware inequality. Widens numeric operands to their common type.-(./=) ::- ( NumericWidenOp (BaseType a) (BaseType b)- , NullLift2Op a b Bool (NullCmpResult a b)- , Eq (Promote (BaseType a) (BaseType b))- ) =>- Expr a ->- Expr b ->- Expr (NullCmpResult a b)-(./=) = lift2Decorated (applyNull2 (widenCmpOp (/=))) "neq" (Just "./=") True 4---- | Nullable-aware less-than. Widens numeric operands to their common type.-(.<) ::- ( NumericWidenOp (BaseType a) (BaseType b)- , NullLift2Op a b Bool (NullCmpResult a b)- , Ord (Promote (BaseType a) (BaseType b))- ) =>- Expr a ->- Expr b ->- Expr (NullCmpResult a b)-(.<) = lift2Decorated (applyNull2 (widenCmpOp (<))) "lt" (Just ".<") False 4---- | Nullable-aware greater-than. Widens numeric operands to their common type.-(.>) ::- ( NumericWidenOp (BaseType a) (BaseType b)- , NullLift2Op a b Bool (NullCmpResult a b)- , Ord (Promote (BaseType a) (BaseType b))- ) =>- Expr a ->- Expr b ->- Expr (NullCmpResult a b)-(.>) = lift2Decorated (applyNull2 (widenCmpOp (>))) "gt" (Just ".>") False 4--{- | Nullable-aware less-than-or-equal. Widens numeric operands to their-common type, so @Expr Double .<= Expr Int@ typechecks.--}-(.<=) ::- ( NumericWidenOp (BaseType a) (BaseType b)- , NullLift2Op a b Bool (NullCmpResult a b)- , Ord (Promote (BaseType a) (BaseType b))- ) =>- Expr a ->- Expr b ->- Expr (NullCmpResult a b)-(.<=) = lift2Decorated (applyNull2 (widenCmpOp (<=))) "leq" (Just ".<=") False 4---- | Nullable-aware greater-than-or-equal. Widens numeric operands to their common type.-(.>=) ::- ( NumericWidenOp (BaseType a) (BaseType b)- , NullLift2Op a b Bool (NullCmpResult a b)- , Ord (Promote (BaseType a) (BaseType b))- ) =>- Expr a ->- Expr b ->- Expr (NullCmpResult a b)-(.>=) = lift2Decorated (applyNull2 (widenCmpOp (>=))) "geq" (Just ".>=") False 4--(.&&.) :: Expr Bool -> Expr Bool -> Expr Bool-(.&&.) = lift2Decorated (&&) "and" (Just ".&&.") True 3--(.||.) :: Expr Bool -> Expr Bool -> Expr Bool-(.||.) = lift2Decorated (||) "or" (Just ".||.") True 2---- | Nullable-aware logical AND. Returns @Maybe Bool@ when either operand is nullable.-(.&&) ::- (NullableCmpOp a b (NullCmpResult a b), BaseType a ~ Bool) =>- Expr a ->- Expr b ->- Expr (NullCmpResult a b)-(.&&) = lift2Decorated (nullCmpOp (&&)) "nulland" (Just ".&&") True 3---- | Nullable-aware logical OR. Returns @Maybe Bool@ when either operand is nullable.-(.||) ::- (NullableCmpOp a b (NullCmpResult a b), BaseType a ~ Bool) =>- Expr a ->- Expr b ->- Expr (NullCmpResult a b)-(.||) = lift2Decorated (nullCmpOp (||)) "nullor" (Just ".||") True 2--(.^^) ::- ( Columnable (BaseType a)- , Columnable (BaseType b)- , Fractional (BaseType a)- , Integral (BaseType b)- , NumericWidenOp (BaseType a) (BaseType b)- , NullLift2Op a b (BaseType a) a- , Num (Promote (BaseType a) (BaseType b))- ) =>- Expr a -> Expr b -> Expr a-(.^^) = lift2Decorated (applyNull2 (^^)) "pow" (Just ".^^") False 8--(.^) ::- ( Columnable (BaseType a)- , Columnable (BaseType b)- , Num (BaseType a)- , Integral (BaseType b)- , NumericWidenOp (BaseType a) (BaseType b)- , NullLift2Op a b (BaseType a) a- , Num (Promote (BaseType a) (BaseType b))- ) =>- Expr a -> Expr b -> Expr a-(.^) = lift2Decorated (applyNull2 (^)) "pow" (Just ".^") False 8---- Same-type (non-nullable) exponentiation operators--(.^^.) ::- (Columnable a, Columnable b, Fractional a, Integral b) =>- Expr a -> Expr b -> Expr a-(.^^.) = lift2Decorated (^^) "pow" (Just ".^^.") False 8--(.^.) ::- (Columnable a, Columnable b, Num a, Integral b) =>- Expr a -> Expr b -> Expr a-(.^.) = lift2Decorated (^) "pow" (Just ".^.") False 8
− src/DataFrame/Synthesis.hs
@@ -1,483 +0,0 @@-{-# LANGUAGE BangPatterns #-}-{-# LANGUAGE ExplicitNamespaces #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE FlexibleInstances #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE MultiParamTypeClasses #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE RankNTypes #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}-{-# LANGUAGE UndecidableInstances #-}--module DataFrame.Synthesis where--import qualified DataFrame.Functions as F-import DataFrame.Internal.Column-import DataFrame.Internal.DataFrame (- DataFrame (..),- )-import DataFrame.Internal.Expression (- Expr (..),- eSize,- eqExpr,- )-import DataFrame.Internal.Interpreter (interpret)-import DataFrame.Internal.Statistics-import DataFrame.Operations.Core (columnAsDoubleVector)-import qualified DataFrame.Operations.Statistics as Stats-import DataFrame.Operations.Subset (exclude)--import Control.Exception (throw)-import Data.Function-import qualified Data.List as L-import qualified Data.Map as M-import Data.Maybe (listToMaybe)-import qualified Data.Text as T-import Data.Type.Equality-import qualified Data.Vector.Unboxed as VU-import qualified DataFrame.Operations.Core as D-import DataFrame.Operators-import Debug.Trace (trace)-import Type.Reflection (typeRep)--generateConditions ::- TypedColumn Double -> [Expr Bool] -> [Expr Double] -> DataFrame -> [Expr Bool]-generateConditions labels conds ps df =- let- newConds =- [ p .<= q- | p <- filter (not . isLiteral) ps- , q <- ps- , Prelude.not (eqExpr p q)- ]- ++ [ F.not p- | p <- conds- ]- expandedConds =- conds- ++ newConds- ++ [p .&& q | p <- newConds, q <- conds, Prelude.not (eqExpr p q)]- ++ [p .|| q | p <- newConds, q <- conds, Prelude.not (eqExpr p q)]- in- pickTopNBool df labels (deduplicate df expandedConds)--generatePrograms ::- Bool ->- [Expr Bool] ->- [Expr Double] ->- [Expr Double] ->- [Expr Double] ->- [Expr Double]-generatePrograms _ _ vars' constants [] = vars' ++ constants-generatePrograms includeConds conds vars constants ps =- let- existingPrograms = ps ++ vars ++ constants- in- existingPrograms- ++ [ transform p- | p <- ps ++ vars- , Prelude.not (isConditional p)- , transform <-- [ sqrt- , abs- , log . (+ Lit 1)- , exp- , sin- , cos- , F.relu- , signum- ]- ]- ++ [ F.pow p i- | p <- existingPrograms- , Prelude.not (isConditional p)- , i <- [2 .. 6]- ]- ++ [ p + q- | (i, p) <- zip [(0 :: Int) ..] existingPrograms- , (j, q) <- zip [(0 :: Int) ..] existingPrograms- , Prelude.not (isLiteral p && isLiteral q)- , Prelude.not (isConditional p || isConditional q)- , i >= j- ]- ++ [ p - q- | (i, p) <- zip [(0 :: Int) ..] existingPrograms- , (j, q) <- zip [(0 :: Int) ..] existingPrograms- , Prelude.not (isLiteral p && isLiteral q)- , Prelude.not (isConditional p || isConditional q)- , i /= j- ]- ++ ( if includeConds- then- [ F.min p q- | (i, p) <- zip [(0 :: Int) ..] existingPrograms- , (j, q) <- zip [(0 :: Int) ..] existingPrograms- , Prelude.not (isLiteral p && isLiteral q)- , Prelude.not (isConditional p || isConditional q)- , Prelude.not (eqExpr p q)- , i > j- ]- ++ [ F.max p q- | (i, p) <- zip [(0 :: Int) ..] existingPrograms- , (j, q) <- zip [(0 :: Int) ..] existingPrograms- , Prelude.not (isLiteral p && isLiteral q)- , Prelude.not (isConditional p || isConditional q)- , Prelude.not (eqExpr p q)- , i > j- ]- ++ [ F.ifThenElse cond r s- | cond <- conds- , r <- existingPrograms- , s <- existingPrograms- , Prelude.not (isConditional r || isConditional s)- , Prelude.not (eqExpr r s)- ]- else []- )- ++ [ p * q- | (i, p) <- zip [(0 :: Int) ..] existingPrograms- , (j, q) <- zip [(0 :: Int) ..] existingPrograms- , Prelude.not (isLiteral p && isLiteral q)- , Prelude.not (isConditional p || isConditional q)- , i >= j- ]- ++ [ p / q- | p <- existingPrograms- , q <- existingPrograms- , Prelude.not (isLiteral p && isLiteral q)- , Prelude.not (isConditional p || isConditional q)- , Prelude.not (eqExpr p q)- ]--isLiteral :: Expr a -> Bool-isLiteral (Lit _) = True-isLiteral _ = False--isConditional :: Expr a -> Bool-isConditional (If{}) = True-isConditional _ = False--deduplicate ::- forall a.- (Columnable a) =>- DataFrame ->- [Expr a] ->- [(Expr a, TypedColumn a)]-deduplicate df = go [] . L.nubBy eqExpr . L.sortBy (\e1 e2 -> compare (eSize e1) (eSize e2))- where- go _ [] = []- go seen (x : xs)- | hasInvalid = go seen xs- | res `elem` seen = go seen xs- | otherwise = (x, res) : go (res : seen) xs- where- res = case interpret @a df x of- Left e -> throw e- Right v -> v- hasInvalid = case res of- (TColumn (UnboxedColumn _ (column :: VU.Vector b))) -> case testEquality (typeRep @Double) (typeRep @b) of- Just Refl -> VU.any (\n -> isNaN n || isInfinite n) column- Nothing -> False- _ -> False---- | Checks if two programs generate the same outputs given all the same inputs.-equivalent :: DataFrame -> Expr Double -> Expr Double -> Bool-equivalent df p1 p2 = case (==) <$> interpret df p1 <*> interpret df p2 of- Left e -> throw e- Right v -> v--synthesizeFeatureExpr ::- -- | Target expression- T.Text ->- BeamConfig ->- DataFrame ->- Either String (Expr Double)-synthesizeFeatureExpr target cfg df =- let- df' = exclude [target] df- t = case interpret df (Col target) of- Left e -> throw e- Right v -> v- in- case beamSearch- df'- cfg- t- (percentiles df')- []- [] of- Nothing -> Left "No programs found"- Just p -> Right p--f1FromBinary :: VU.Vector Double -> VU.Vector Double -> Maybe Double-f1FromBinary trues preds =- let (!tp, !fp, !fn) =- VU.foldl' step (0 :: Int, 0 :: Int, 0 :: Int) $- VU.zip (VU.map (> 0) preds) (VU.map (> 0) trues)- in f1FromCounts tp fp fn- where- step (!tp, !fp, !fn) (!p, !t) =- case (p, t) of- (True, True) -> (tp + 1, fp, fn)- (True, False) -> (tp, fp + 1, fn)- (False, True) -> (tp, fp, fn + 1)- (False, False) -> (tp, fp, fn)--f1FromCounts :: Int -> Int -> Int -> Maybe Double-f1FromCounts tp fp fn =- let tp' = fromIntegral tp- fp' = fromIntegral fp- fn' = fromIntegral fn- precision = if tp' + fp' == 0 then 0 else tp' / (tp' + fp')- recall = if tp' + fn' == 0 then 0 else tp' / (tp' + fn')- in if precision + recall == 0- then Nothing- else Just (2 * precision * recall / (precision + recall))--fitClassifier ::- -- | Target expression- T.Text ->- -- | Depth of search (Roughly, how many terms in the final expression)- Int ->- -- | Beam size - the number of candidate expressions to consider at a time.- Int ->- DataFrame ->- Either String (Expr Int)-fitClassifier target d b df =- let- df' = exclude [target] df- t = case interpret df (Col target) of- Left e -> throw e- Right v -> v- in- case beamSearch- df'- (BeamConfig d b F1 True)- t- (percentiles df' ++ [Lit 1, Lit 0, Lit (-1)])- []- [] of- Nothing -> Left "No programs found"- Just p -> Right (F.ifThenElse (p .> (0 :: Expr Double)) 1 0)--percentiles :: DataFrame -> [Expr Double]-percentiles df =- let- doubleColumns =- map- (either throw id . ((`columnAsDoubleVector` df) . Col @Double))- (D.columnNames df)- in- concatMap- (\c -> map (Lit . roundTo2SigDigits . (`percentile'` c)) [1, 25, 75, 99])- doubleColumns- ++ map (Lit . roundTo2SigDigits . variance') doubleColumns- ++ map (Lit . roundTo2SigDigits . sqrt . variance') doubleColumns--roundToSigDigits :: Int -> Double -> Double-roundToSigDigits n x- | x == 0 = 0- | otherwise =- let magnitude = floor (logBase 10 (abs x))- scale = 10 ** fromIntegral (n - 1 - magnitude)- in fromIntegral (round (x * scale) :: Int) / scale--roundTo2SigDigits :: Double -> Double-roundTo2SigDigits = roundToSigDigits 2--fitRegression ::- -- | Target expression- T.Text ->- -- | Depth of search (Roughly, how many terms in the final expression)- Int ->- -- | Beam size - the number of candidate expressions to consider at a time.- Int ->- DataFrame ->- Either String (Expr Double)-fitRegression target d b df =- let- df' = exclude [target] df- targetMean = Stats.mean (Col @Double target) df- t = case interpret df (Col target) of- Left e -> throw e- Right v -> v- cfg = BeamConfig d b MeanSquaredError True- constants =- percentiles df'- ++ [Lit targetMean]- ++ [ F.pow p i- | i <- [1 .. 6]- , p <- [Lit 10, Lit 1, Lit 0.1]- ]- in- case beamSearch df' cfg t constants [] [] of- Nothing -> Left "No programs found"- Just p -> Right p--data LossFunction- = PearsonCorrelation- | MutualInformation- | MeanSquaredError- | F1--getLossFunction ::- LossFunction -> (VU.Vector Double -> VU.Vector Double -> Maybe Double)-getLossFunction f = case f of- MutualInformation ->- ( \l r ->- mutualInformationBinned- (Prelude.max 10 (ceiling (sqrt (fromIntegral (VU.length l) :: Double))))- l- r- )- PearsonCorrelation -> (\l r -> (^ (2 :: Int)) <$> correlation' l r)- MeanSquaredError -> (\l r -> fmap negate (meanSquaredError l r))- F1 -> f1FromBinary--data BeamConfig = BeamConfig- { searchDepth :: Int- , beamLength :: Int- , lossFunction :: LossFunction- , includeConditionals :: Bool- }--defaultBeamConfig :: BeamConfig-defaultBeamConfig = BeamConfig 2 100 PearsonCorrelation False--beamSearch ::- DataFrame ->- -- | Parameters of the beam search.- BeamConfig ->- -- | Examples- TypedColumn Double ->- -- | Constants- [Expr Double] ->- -- | Conditions- [Expr Bool] ->- -- | Programs- [Expr Double] ->- Maybe (Expr Double)-beamSearch df cfg outputs constants conds programs- | searchDepth cfg == 0 = case ps of- [] -> Nothing- (x : _) -> Just x- | otherwise =- beamSearch- df- (cfg{searchDepth = searchDepth cfg - 1})- outputs- constants- conditions- (generatePrograms (includeConditionals cfg) conditions vars constants ps)- where- vars = map Col names- conditions = generateConditions outputs conds (vars ++ constants) df- ps = pickTopN df outputs cfg $ deduplicate df programs- names = (map fst . L.sortBy (compare `on` snd) . M.toList . columnIndices) df--pickTopN ::- DataFrame ->- TypedColumn Double ->- BeamConfig ->- [(Expr Double, TypedColumn a)] ->- [Expr Double]-pickTopN _ _ _ [] = []-pickTopN df (TColumn column) cfg ps =- let- l = case toVector @Double @VU.Vector column of- Left e -> throw e- Right v -> v- ordered =- Prelude.take- (beamLength cfg)- ( map fst $- L.sortBy- ( \(_, c2) (_, c1) ->- if maybe False isInfinite c1- || maybe False isInfinite c2- || maybe False isNaN c1- || maybe False isNaN c2- then LT- else compare c1 c2- )- ( map- (\(e, res) -> (e, getLossFunction (lossFunction cfg) l (asDoubleVector res)))- ps- )- )- asDoubleVector c =- let- (TColumn col') = c- in- case toVector @Double @VU.Vector col' of- Left e -> throw e- Right v -> VU.convert v- interpretDoubleVector e' =- let- (TColumn col') = case interpret df e' of- Left err -> throw err- Right v -> v- in- case toVector @Double @VU.Vector col' of- Left err -> throw err- Right v -> VU.convert v- in- trace- ( "Best loss: "- ++ show- ( getLossFunction (lossFunction cfg) l . interpretDoubleVector- <$> listToMaybe ordered- )- ++ " "- ++ (if null ordered then "empty" else show (listToMaybe ordered))- )- ordered--pickTopNBool ::- DataFrame ->- TypedColumn Double ->- [(Expr Bool, TypedColumn Bool)] ->- [Expr Bool]-pickTopNBool _ _ [] = []-pickTopNBool _df (TColumn column) ps =- let- l = case toVector @Double @VU.Vector column of- Left e -> throw e- Right v -> v- ordered =- Prelude.take- 10- ( map fst $- L.sortBy- ( \(_, c2) (_, c1) ->- if maybe False isInfinite c1- || maybe False isInfinite c2- || maybe False isNaN c1- || maybe False isNaN c2- then LT- else compare c1 c2- )- ( map- (\(e, res) -> (e, getLossFunction MutualInformation l (asDoubleVector res)))- ps- )- )- asDoubleVector c =- let- (TColumn col') = c- in- case toVector @Bool @VU.Vector col' of- Left e -> throw e- Right v -> VU.map (fromIntegral @Int @Double . fromEnum) v- in- ordered--satisfiesExamples :: DataFrame -> TypedColumn Double -> Expr Double -> Bool-satisfiesExamples df column expr =- let- result = case interpret df expr of- Left e -> throw e- Right v -> v- in- result == column
src/DataFrame/TH.hs view
@@ -1,213 +1,35 @@-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TemplateHaskell #-}-{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE CPP #-} {- | 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@.+Backwards-compatibility re-export hub for the split @DataFrame.TH.*@+modules. New code can import the specific submodule directly: -@-ghci> :set -XTemplateHaskell-ghci> :declareColumns df-ghci> :type passengers-passengers :: Expr Int-@+ * "DataFrame.TH.Records" — record/DataFrame-based splices (no file IO)+ * "DataFrame.TH.CSV" — CSV-file-based splices (requires @-fwith-csv@)+ * "DataFrame.TH.Parquet" — Parquet-file-based splices (requires @-fwith-parquet@) -The typed-API equivalents (which generate a schema type synonym) live in-"DataFrame.Typed.TH".+These live in @dataframe-th@, @dataframe-csv-th@, and+@dataframe-parquet-th@ respectively. The CSV/Parquet re-exports are+guarded with the @CSV_TH@ / @PARQUET_TH@ CPP defines that the+meta-@dataframe@ package sets based on its cabal flags. -} module DataFrame.TH (- -- * Declare one binding per column- declareColumns,- declareColumnsWithPrefix,- declareColumnsWithPrefix',-- -- * From a file- declareColumnsFromCsvFile,- declareColumnsFromCsvWithOpts,- declareColumnsFromParquetFile,-- -- * Type-string parser (exposed for testing)- typeFromString,+ module DataFrame.TH.Records,+#ifdef WITH_CSV_TH+ module DataFrame.TH.CSV,+#endif+#ifdef WITH_PARQUET_TH+ module DataFrame.TH.Parquet,+#endif ) 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]+import DataFrame.TH.Records+#ifdef WITH_CSV_TH+import DataFrame.TH.CSV+#endif+#ifdef WITH_PARQUET_TH+import DataFrame.TH.Parquet+#endif
src/DataFrame/Typed.hs view
@@ -1,3 +1,4 @@+{-# LANGUAGE CPP #-} {-# LANGUAGE DataKinds #-} {- |@@ -185,14 +186,18 @@ aggregate, aggregateUntyped, +#ifdef WITH_TH -- * Template Haskell deriveSchema,+#ifdef WITH_CSV_TH deriveSchemaFromCsvFile, deriveSchemaFromCsvFileWith,+#endif deriveSchemaFromType, deriveSchemaFromTypeWith, SchemaOptions (..), defaultSchemaOptions,+#endif -- * Record bridge (ADT <-> TypedDataFrame) HasSchema (..),@@ -260,15 +265,19 @@ toRecordsTyped, ) import DataFrame.Typed.Schema+#ifdef WITH_TH import DataFrame.Typed.TH ( SchemaOptions (..), defaultSchemaOptions, deriveSchema,+#ifdef WITH_CSV_TH deriveSchemaFromCsvFile, deriveSchemaFromCsvFileWith,+#endif deriveSchemaFromType, deriveSchemaFromTypeWith, )+#endif import DataFrame.Typed.Types ( Column, TSortOrder (..),
− src/DataFrame/Typed/Access.hs
@@ -1,55 +0,0 @@-{-# LANGUAGE AllowAmbiguousTypes #-}-{-# LANGUAGE DataKinds #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}-{-# LANGUAGE TypeFamilies #-}-{-# LANGUAGE TypeOperators #-}--module DataFrame.Typed.Access (- -- * Typed column access- columnAsVector,- columnAsList,-) where--import Control.Exception (throw)-import Data.Proxy (Proxy (..))-import qualified Data.Text as T-import qualified Data.Vector as V-import GHC.TypeLits (KnownSymbol, symbolVal)--import DataFrame.Internal.Column (Columnable)-import DataFrame.Internal.Expression (Expr (Col))-import qualified DataFrame.Operations.Core as D-import DataFrame.Typed.Schema (AssertPresent, SafeLookup)-import DataFrame.Typed.Types (TypedDataFrame (..))--{- | Retrieve a column as a boxed 'Vector', with the type determined by-the schema. The column must exist (enforced at compile time).--}-columnAsVector ::- forall name cols a.- ( KnownSymbol name- , a ~ SafeLookup name cols- , Columnable a- , AssertPresent name cols- ) =>- TypedDataFrame cols -> V.Vector a-columnAsVector (TDF df) =- either throw id $ D.columnAsVector (Col @a colName) df- where- colName = T.pack (symbolVal (Proxy @name))---- | Retrieve a column as a list, with the type determined by the schema.-columnAsList ::- forall name cols a.- ( KnownSymbol name- , a ~ SafeLookup name cols- , Columnable a- , AssertPresent name cols- ) =>- TypedDataFrame cols -> [a]-columnAsList (TDF df) =- D.columnAsList (Col @a colName) df- where- colName = T.pack (symbolVal (Proxy @name))
− src/DataFrame/Typed/Aggregate.hs
@@ -1,118 +0,0 @@-{-# LANGUAGE AllowAmbiguousTypes #-}-{-# LANGUAGE DataKinds #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}-{-# LANGUAGE TypeFamilies #-}-{-# LANGUAGE TypeOperators #-}--module DataFrame.Typed.Aggregate (- -- * Typed groupBy- groupBy,-- -- * Naming an aggregation- as,-- -- * Running aggregations- aggregate,-- -- * Escape hatch- aggregateUntyped,-) where--import Data.Proxy (Proxy (..))-import qualified Data.Text as T-import GHC.TypeLits (KnownSymbol, Symbol, symbolVal)--import DataFrame.Internal.Column (Columnable)-import qualified DataFrame.Internal.DataFrame as D-import DataFrame.Internal.Expression (NamedExpr)-import qualified DataFrame.Operations.Aggregation as DA--import DataFrame.Typed.Freeze (unsafeFreeze)-import DataFrame.Typed.Schema-import DataFrame.Typed.Types--{- | Group a typed DataFrame by one or more key columns.--@-grouped = groupBy \@'[\"department\"] employees-@--}-groupBy ::- forall (keys :: [Symbol]) cols.- (AllKnownSymbol keys, AssertAllPresent keys cols) =>- TypedDataFrame cols -> TypedGrouped keys cols-groupBy (TDF df) = TGD (DA.groupBy (symbolVals @keys) df)--{- | 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.--@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.--==== __Prefix form__--@-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 @(|>)@.--}-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)-as = TAggCons (T.pack (symbolVal (Proxy @name)))--{- | Run a typed aggregation against a grouped DataFrame.--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 = 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 '[] -> TAgg keys cols aggs) ->- TypedGrouped keys cols ->- TypedDataFrame (Append (GroupKeyColumns keys cols) (Reverse aggs))-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-aggregateUntyped exprs (TGD gdf) = DA.aggregate exprs gdf
− src/DataFrame/Typed/Expr.hs
@@ -1,644 +0,0 @@-{-# LANGUAGE AllowAmbiguousTypes #-}-{-# LANGUAGE DataKinds #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE FlexibleInstances #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE MultiParamTypeClasses #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE RankNTypes #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}-{-# LANGUAGE TypeFamilies #-}-{-# LANGUAGE TypeOperators #-}-{-# LANGUAGE UndecidableInstances #-}-{-# OPTIONS_GHC -Wno-orphans #-}--{- | Type-safe expression construction for typed DataFrames.--Unlike the untyped @Expr a@ where column references are unchecked strings,-'TExpr' ensures at compile time that:--* Referenced columns exist in the schema-* Column types match the expression type--== Example--@-type Schema = '[Column \"age\" Int, Column \"salary\" Double]---- This compiles:-goodExpr :: TExpr Schema Double-goodExpr = col \@\"salary\"---- This gives a compile-time error (column not found):-badExpr :: TExpr Schema Double-badExpr = col \@\"nonexistent\"---- This gives a compile-time error (type mismatch):-wrongType :: TExpr Schema Int-wrongType = col \@\"salary\" -- salary is Double, not Int-@--}-module DataFrame.Typed.Expr (- -- * Core typed expression type (re-exported from Types)- TExpr (..),-- -- * Column reference (schema-checked)- col,-- -- * Literals- lit,-- -- * Conditional- ifThenElse,-- -- * Unary / binary lifting- lift,- lift2,- nullLift,- nullLift2,-- -- * Same-type comparison operators- (.==.),- (./=.),- (.<.),- (.<=.),- (.>=.),- (.>.),-- -- * Same-type arithmetic operators- (.+.),- (.-.),- (.*.),- (./.),-- -- * Same-type exponentiation operators- (.^^.),- (.^.),-- -- * Nullable-aware arithmetic operators- (.+),- (.-),- (.*),- (./),-- -- * Nullable-aware exponentiation operators- (.^^),- (.^),-- -- * Nullable-aware comparison operators (three-valued logic)- (.==),- (./=),- (.<),- (.<=),- (.>=),- (.>),-- -- * Logical operators- (.&&.),- (.||.),- (.&&),- (.||),- DataFrame.Typed.Expr.not,-- -- * Aggregation combinators- sum,- mean,- median,- count,- countAll,- minimum,- maximum,- collect,- over,-- -- * Cast / coercion expressions- castExpr,- castExprWithDefault,- castExprEither,- unsafeCastExpr,- toDouble,-- -- * Sort helpers- asc,- desc,-) where--import Data.Either (fromRight)-import Data.Proxy (Proxy (..))-import Data.String (IsString (..))-import qualified Data.Text as T-import GHC.TypeLits (KnownSymbol, Symbol, symbolVal)--import qualified DataFrame.Functions as F-import DataFrame.Internal.Column (Columnable)-import DataFrame.Internal.Expression (- BinaryOp (..),- Expr (..),- UnaryOp (..),- )-import DataFrame.Internal.Nullable (- BaseType,- DivWidenOp,- NullCmpResult,- NullLift1Op (applyNull1),- NullLift1Result,- NullLift2Op (applyNull2),- NullLift2Result,- NullableCmpOp (nullCmpOp),- NumericWidenOp,- WidenResult,- WidenResultDiv,- divArithOp,- widenArithOp,- widenCmpOp,- )-import DataFrame.Internal.Types (Promote, PromoteDiv)--import qualified Data.Vector.Unboxed as VU-import DataFrame.Typed.Schema (- AllKnownSymbol,- AssertAllPresent,- AssertPresent,- SafeLookup,- symbolVals,- )-import DataFrame.Typed.Types (TExpr (..), TSortOrder (..))-import Prelude hiding (maximum, minimum, sum)--{- | Create a typed column reference. This is the key type-safety entry point.--The column name must exist in @cols@ and its type must match @a@.-Both checks happen at compile time via type families.--@-salary :: TExpr '[Column \"salary\" Double] Double-salary = col \@\"salary\"-@--}-col ::- forall (name :: Symbol) cols a.- ( KnownSymbol name- , a ~ SafeLookup name cols- , Columnable a- , AssertPresent name cols- ) =>- TExpr cols a-col = TExpr (Col (T.pack (symbolVal (Proxy @name))))--{- | Create a literal expression. Valid for any schema since it-references no columns.--}-lit :: (Columnable a) => a -> TExpr cols a-lit = TExpr . Lit---- | Conditional expression.-ifThenElse ::- (Columnable a) =>- TExpr cols Bool -> TExpr cols a -> TExpr cols a -> TExpr cols a-ifThenElse (TExpr c) (TExpr t) (TExpr e) = TExpr (If c t e)------------------------------------------------------------------------------------ Numeric instances (mirror Expr's instances)----------------------------------------------------------------------------------instance (Num a, Columnable a) => Num (TExpr cols a) where- (TExpr a) + (TExpr b) = TExpr (a + b)- (TExpr a) - (TExpr b) = TExpr (a - b)- (TExpr a) * (TExpr b) = TExpr (a * b)- negate (TExpr a) = TExpr (negate a)- abs (TExpr a) = TExpr (abs a)- signum (TExpr a) = TExpr (signum a)- fromInteger = TExpr . fromInteger--instance (Fractional a, Columnable a) => Fractional (TExpr cols a) where- fromRational = TExpr . fromRational- (TExpr a) / (TExpr b) = TExpr (a / b)--instance (Floating a, Columnable a) => Floating (TExpr cols a) where- pi = TExpr pi- exp (TExpr a) = TExpr (exp a)- sqrt (TExpr a) = TExpr (sqrt a)- log (TExpr a) = TExpr (log a)- (TExpr a) ** (TExpr b) = TExpr (a ** b)- logBase (TExpr a) (TExpr b) = TExpr (logBase a b)- sin (TExpr a) = TExpr (sin a)- cos (TExpr a) = TExpr (cos a)- tan (TExpr a) = TExpr (tan a)- asin (TExpr a) = TExpr (asin a)- acos (TExpr a) = TExpr (acos a)- atan (TExpr a) = TExpr (atan a)- sinh (TExpr a) = TExpr (sinh a)- cosh (TExpr a) = TExpr (cosh a)- asinh (TExpr a) = TExpr (asinh a)- acosh (TExpr a) = TExpr (acosh a)- atanh (TExpr a) = TExpr (atanh a)--instance (IsString a, Columnable a) => IsString (TExpr cols a) where- fromString = TExpr . fromString------------------------------------------------------------------------------------ Lifting arbitrary functions------------------------------------------------------------------------------------ | Lift a unary function into a typed expression.-lift ::- (Columnable a, Columnable b) => (a -> b) -> TExpr cols a -> TExpr cols b-lift f (TExpr e) = TExpr (Unary (MkUnaryOp f "unaryUdf" Nothing) e)---- | Lift a binary function into typed expressions.-lift2 ::- (Columnable a, Columnable b, Columnable c) =>- (a -> b -> c) -> TExpr cols a -> TExpr cols b -> TExpr cols c-lift2 f (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp f "binaryUdf" Nothing False 0) a b)--{- | Typed 'nullLift': lift a unary function with nullable propagation.-When the input is @Maybe a@, 'Nothing' short-circuits; when plain @a@, applies directly.-The return type is inferred via 'NullLift1Result': no annotation needed.--}-nullLift ::- (NullLift1Op a r (NullLift1Result a r), Columnable (NullLift1Result a r)) =>- (BaseType a -> r) ->- TExpr cols a ->- TExpr cols (NullLift1Result a r)-nullLift f (TExpr e) = TExpr (Unary (MkUnaryOp (applyNull1 f) "nullLift" Nothing) e)--{- | Typed 'nullLift2': lift a binary function with nullable propagation.-Any 'Nothing' operand short-circuits to 'Nothing' in the result.-The return type is inferred via 'NullLift2Result': no annotation needed.--}-nullLift2 ::- (NullLift2Op a b r (NullLift2Result a b r), Columnable (NullLift2Result a b r)) =>- (BaseType a -> BaseType b -> r) ->- TExpr cols a ->- TExpr cols b ->- TExpr cols (NullLift2Result a b r)-nullLift2 f (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (applyNull2 f) "nullLift2" Nothing False 0) a b)--infixl 4 .==., ./=., .<., .<=., .>=., .>.-infix 4 .==, ./=, .<, .<=, .>=, .>-infixr 3 .&&., .&&-infixr 2 .||., .||-infixl 6 .+., .-.-infixl 7 .*., ./.-infixr 8 .^^., .^^, .^., .^--(.==.) ::- (Columnable a, Eq a) => TExpr cols a -> TExpr cols a -> TExpr cols Bool-(.==.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (==) "eq" (Just "==") True 4) a b)--(./=.) ::- (Columnable a, Eq a) => TExpr cols a -> TExpr cols a -> TExpr cols Bool-(./=.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (/=) "neq" (Just "/=") True 4) a b)--(.<.) ::- (Columnable a, Ord a) => TExpr cols a -> TExpr cols a -> TExpr cols Bool-(.<.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (<) "lt" (Just "<") False 4) a b)--(.<=.) ::- (Columnable a, Ord a) => TExpr cols a -> TExpr cols a -> TExpr cols Bool-(.<=.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (<=) "leq" (Just "<=") False 4) a b)--(.>=.) ::- (Columnable a, Ord a) => TExpr cols a -> TExpr cols a -> TExpr cols Bool-(.>=.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (>=) "geq" (Just ">=") False 4) a b)--(.>.) ::- (Columnable a, Ord a) => TExpr cols a -> TExpr cols a -> TExpr cols Bool-(.>.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (>) "gt" (Just ">") False 4) a b)---- Same-type arithmetic operators--(.+.) :: (Columnable a, Num a) => TExpr cols a -> TExpr cols a -> TExpr cols a-(.+.) = (+)--(.-.) :: (Columnable a, Num a) => TExpr cols a -> TExpr cols a -> TExpr cols a-(.-.) = (-)--(.*.) :: (Columnable a, Num a) => TExpr cols a -> TExpr cols a -> TExpr cols a-(.*.) = (*)--(./.) ::- (Columnable a, Fractional a) => TExpr cols a -> TExpr cols a -> TExpr cols a-(./.) = (/)---- Same-type exponentiation operators--(.^^.) ::- (Columnable a, Columnable b, Fractional a, Integral b) =>- TExpr cols a -> TExpr cols b -> TExpr cols a-(.^^.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (^^) "pow" (Just ".^^.") False 8) a b)--(.^.) ::- (Columnable a, Columnable b, Num a, Integral b) =>- TExpr cols a -> TExpr cols b -> TExpr cols a-(.^.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (^) "pow" (Just ".^.") False 8) a b)--(.&&.) :: TExpr cols Bool -> TExpr cols Bool -> TExpr cols Bool-(.&&.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (&&) "and" (Just ".&&.") True 3) a b)--(.||.) :: TExpr cols Bool -> TExpr cols Bool -> TExpr cols Bool-(.||.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (||) "or" (Just ".||.") True 2) a b)---- | Nullable-aware logical AND. Returns @Maybe Bool@ when either operand is nullable.-(.&&) ::- (NullableCmpOp a b (NullCmpResult a b), BaseType a ~ Bool) =>- TExpr cols a ->- TExpr cols b ->- TExpr cols (NullCmpResult a b)-(.&&) (TExpr a) (TExpr b) =- TExpr (Binary (MkBinaryOp (nullCmpOp (&&)) "nulland" (Just ".&&") True 3) a b)---- | Nullable-aware logical OR. Returns @Maybe Bool@ when either operand is nullable.-(.||) ::- (NullableCmpOp a b (NullCmpResult a b), BaseType a ~ Bool) =>- TExpr cols a ->- TExpr cols b ->- TExpr cols (NullCmpResult a b)-(.||) (TExpr a) (TExpr b) =- TExpr (Binary (MkBinaryOp (nullCmpOp (||)) "nullor" (Just ".||") True 2) a b)------------------------------------------------------------------------------------ Nullable-aware arithmetic operators----------------------------------------------------------------------------------infixl 6 .+, .--infixl 7 .*, ./--{- | Nullable-aware addition. Works for all combinations of nullable\/non-nullable operands.-@col \@\"x\" '.+' col \@\"y\" -- :: TExpr cols (Maybe Int) when y :: Maybe Int@--}-(.+) ::- ( NumericWidenOp (BaseType a) (BaseType b)- , NullLift2Op a b (Promote (BaseType a) (BaseType b)) (WidenResult a b)- , Num (Promote (BaseType a) (BaseType b))- ) =>- TExpr cols a ->- TExpr cols b ->- TExpr cols (WidenResult a b)-(.+) (TExpr a) (TExpr b) =- TExpr- ( Binary- (MkBinaryOp (applyNull2 (widenArithOp (+))) "nulladd" (Just "+") True 6)- a- b- )---- | Nullable-aware subtraction.-(.-) ::- ( NumericWidenOp (BaseType a) (BaseType b)- , NullLift2Op a b (Promote (BaseType a) (BaseType b)) (WidenResult a b)- , Num (Promote (BaseType a) (BaseType b))- ) =>- TExpr cols a ->- TExpr cols b ->- TExpr cols (WidenResult a b)-(.-) (TExpr a) (TExpr b) =- TExpr- ( Binary- (MkBinaryOp (applyNull2 (widenArithOp (-))) "nullsub" (Just "-") False 6)- a- b- )---- | Nullable-aware multiplication.-(.*) ::- ( NumericWidenOp (BaseType a) (BaseType b)- , NullLift2Op a b (Promote (BaseType a) (BaseType b)) (WidenResult a b)- , Num (Promote (BaseType a) (BaseType b))- ) =>- TExpr cols a ->- TExpr cols b ->- TExpr cols (WidenResult a b)-(.*) (TExpr a) (TExpr b) =- TExpr- ( Binary- (MkBinaryOp (applyNull2 (widenArithOp (*))) "nullmul" (Just "*") True 7)- a- b- )---- | Nullable-aware division. Integral operands are promoted to Double.-(./) ::- ( DivWidenOp (BaseType a) (BaseType b)- , NullLift2Op a b (PromoteDiv (BaseType a) (BaseType b)) (WidenResultDiv a b)- , Fractional (PromoteDiv (BaseType a) (BaseType b))- ) =>- TExpr cols a ->- TExpr cols b ->- TExpr cols (WidenResultDiv a b)-(./) (TExpr a) (TExpr b) =- TExpr- ( Binary- (MkBinaryOp (applyNull2 (divArithOp (/))) "nulldiv" (Just "/") False 7)- a- b- )---- | Nullable-aware exponentiation (fractional base, integral exponent).-(.^^) ::- ( Columnable (BaseType a)- , Columnable (BaseType b)- , Fractional (BaseType a)- , Integral (BaseType b)- , NumericWidenOp (BaseType a) (BaseType b)- , NullLift2Op a b (BaseType a) a- , Num (Promote (BaseType a) (BaseType b))- ) =>- TExpr cols a -> TExpr cols b -> TExpr cols a-(.^^) (TExpr a) (TExpr b) =- TExpr (Binary (MkBinaryOp (applyNull2 (^^)) "pow" (Just ".^^") False 8) a b)---- | Nullable-aware exponentiation (num base, integral exponent).-(.^) ::- ( Columnable (BaseType a)- , Columnable (BaseType b)- , Num (BaseType a)- , Integral (BaseType b)- , NumericWidenOp (BaseType a) (BaseType b)- , NullLift2Op a b (BaseType a) a- , Num (Promote (BaseType a) (BaseType b))- ) =>- TExpr cols a -> TExpr cols b -> TExpr cols a-(.^) (TExpr a) (TExpr b) =- TExpr (Binary (MkBinaryOp (applyNull2 (^)) "pow" (Just ".^") False 8) a b)------------------------------------------------------------------------------------ Nullable-aware comparison operators (three-valued logic)----------------------------------------------------------------------------------{- | Nullable-aware equality. Widens numeric operands to their common type,-so @TExpr cols Double .== TExpr cols Int@ typechecks. Returns @Maybe Bool@-when either operand is nullable.--}-(.==) ::- ( NumericWidenOp (BaseType a) (BaseType b)- , NullLift2Op a b Bool (NullCmpResult a b)- , Eq (Promote (BaseType a) (BaseType b))- ) =>- TExpr cols a ->- TExpr cols b ->- TExpr cols (NullCmpResult a b)-(.==) (TExpr a) (TExpr b) =- TExpr- (Binary (MkBinaryOp (applyNull2 (widenCmpOp (==))) "eq" (Just "==") True 4) a b)---- | Nullable-aware inequality. Widens numeric operands to their common type.-(./=) ::- ( NumericWidenOp (BaseType a) (BaseType b)- , NullLift2Op a b Bool (NullCmpResult a b)- , Eq (Promote (BaseType a) (BaseType b))- ) =>- TExpr cols a ->- TExpr cols b ->- TExpr cols (NullCmpResult a b)-(./=) (TExpr a) (TExpr b) =- TExpr- (Binary (MkBinaryOp (applyNull2 (widenCmpOp (/=))) "neq" (Just "/=") True 4) a b)---- | Nullable-aware less-than. Widens numeric operands to their common type.-(.<) ::- ( NumericWidenOp (BaseType a) (BaseType b)- , NullLift2Op a b Bool (NullCmpResult a b)- , Ord (Promote (BaseType a) (BaseType b))- ) =>- TExpr cols a ->- TExpr cols b ->- TExpr cols (NullCmpResult a b)-(.<) (TExpr a) (TExpr b) =- TExpr- (Binary (MkBinaryOp (applyNull2 (widenCmpOp (<))) "lt" (Just "<") False 4) a b)---- | Nullable-aware less-than-or-equal. Widens numeric operands to their common type.-(.<=) ::- ( NumericWidenOp (BaseType a) (BaseType b)- , NullLift2Op a b Bool (NullCmpResult a b)- , Ord (Promote (BaseType a) (BaseType b))- ) =>- TExpr cols a ->- TExpr cols b ->- TExpr cols (NullCmpResult a b)-(.<=) (TExpr a) (TExpr b) =- TExpr- (Binary (MkBinaryOp (applyNull2 (widenCmpOp (<=))) "leq" (Just "<=") False 4) a b)---- | Nullable-aware greater-than-or-equal. Widens numeric operands to their common type.-(.>=) ::- ( NumericWidenOp (BaseType a) (BaseType b)- , NullLift2Op a b Bool (NullCmpResult a b)- , Ord (Promote (BaseType a) (BaseType b))- ) =>- TExpr cols a ->- TExpr cols b ->- TExpr cols (NullCmpResult a b)-(.>=) (TExpr a) (TExpr b) =- TExpr- (Binary (MkBinaryOp (applyNull2 (widenCmpOp (>=))) "geq" (Just ">=") False 4) a b)---- | Nullable-aware greater-than. Widens numeric operands to their common type.-(.>) ::- ( NumericWidenOp (BaseType a) (BaseType b)- , NullLift2Op a b Bool (NullCmpResult a b)- , Ord (Promote (BaseType a) (BaseType b))- ) =>- TExpr cols a ->- TExpr cols b ->- TExpr cols (NullCmpResult a b)-(.>) (TExpr a) (TExpr b) =- TExpr- (Binary (MkBinaryOp (applyNull2 (widenCmpOp (>))) "gt" (Just ">") False 4) a b)--not :: TExpr cols Bool -> TExpr cols Bool-not (TExpr e) = TExpr (Unary (MkUnaryOp Prelude.not "not" (Just "!")) e)------------------------------------------------------------------------------------ Aggregation combinators----------------------------------------------------------------------------------sum :: (Columnable a, Num a) => TExpr cols a -> TExpr cols a-sum (TExpr e) = TExpr (F.sum e)--mean :: (Columnable a, Real a) => TExpr cols a -> TExpr cols Double-mean (TExpr e) = TExpr (F.mean e)--median ::- (Columnable a, Real a, VU.Unbox a) => TExpr cols a -> TExpr cols Double-median (TExpr e) = TExpr (F.median e)--count :: (Columnable a) => TExpr cols a -> TExpr cols Int-count (TExpr e) = TExpr (F.count e)---- | Row count, the equivalent of SQL's @COUNT(*)@.-countAll :: TExpr cols Int-countAll = TExpr F.countAll--minimum :: (Columnable a, Ord a) => TExpr cols a -> TExpr cols a-minimum (TExpr e) = TExpr (F.minimum e)--maximum :: (Columnable a, Ord a) => TExpr cols a -> TExpr cols a-maximum (TExpr e) = TExpr (F.maximum e)--collect :: (Columnable a) => TExpr cols a -> TExpr cols [a]-collect (TExpr e) = TExpr (F.collect e)--over ::- forall (names :: [Symbol]) cols a.- (Columnable a, AllKnownSymbol names, AssertAllPresent names cols) =>- TExpr cols a -> TExpr cols a-over (TExpr e) = TExpr{unTExpr = F.over (symbolVals @names) e}------------------------------------------------------------------------------------ Cast / coercion expressions----------------------------------------------------------------------------------castExpr ::- forall b cols src.- (Columnable b, Columnable src, Read b) => TExpr cols src -> TExpr cols (Maybe b)-castExpr (TExpr e) =- TExpr- (CastExprWith @b @(Maybe b) @src "castExpr" (either (const Nothing) Just) e)--castExprWithDefault ::- forall b cols src.- (Columnable b, Columnable src, Read b) => b -> TExpr cols src -> TExpr cols b-castExprWithDefault def (TExpr e) =- TExpr- ( CastExprWith @b @b @src- ("castExprWithDefault:" <> T.pack (show def))- (fromRight def)- e- )--castExprEither ::- forall b cols src.- (Columnable b, Columnable src, Read b) =>- TExpr cols src -> TExpr cols (Either T.Text b)-castExprEither (TExpr e) =- TExpr- ( CastExprWith @b @(Either T.Text b) @src- "castExprEither"- (either (Left . T.pack) Right)- e- )--unsafeCastExpr ::- forall b cols src.- (Columnable b, Columnable src, Read b) => TExpr cols src -> TExpr cols b-unsafeCastExpr (TExpr e) =- TExpr- ( CastExprWith @b @b @src- "unsafeCastExpr"- (fromRight (error "unsafeCastExpr: unexpected Nothing in column"))- e- )--toDouble :: (Columnable a, Real a) => TExpr cols a -> TExpr cols Double-toDouble (TExpr e) = TExpr (F.toDouble e)---- | Create an ascending sort order from a typed expression.-asc :: (Columnable a, Ord a) => TExpr cols a -> TSortOrder cols-asc = Asc---- | Create a descending sort order from a typed expression.-desc :: (Columnable a, Ord a) => TExpr cols a -> TSortOrder cols-desc = Desc
− src/DataFrame/Typed/Freeze.hs
@@ -1,98 +0,0 @@-{-# LANGUAGE AllowAmbiguousTypes #-}-{-# LANGUAGE DataKinds #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}--module DataFrame.Typed.Freeze (- -- * Safe boundary- freeze,- freezeWithError,-- -- * Escape hatches- thaw,- unsafeFreeze,-) where--import qualified Data.Text as T-import Type.Reflection (SomeTypeRep)--import Data.List (stripPrefix)-import qualified DataFrame.Internal.Column as C-import qualified DataFrame.Internal.DataFrame as D-import DataFrame.Operations.Core (columnNames)-import DataFrame.Typed.Schema (KnownSchema (..))-import DataFrame.Typed.Types (TypedDataFrame (..))--{- | Validate that an untyped 'DataFrame' matches the expected schema @cols@,-then wrap it. Returns 'Nothing' on mismatch.--}-freeze ::- forall cols. (KnownSchema cols) => D.DataFrame -> Maybe (TypedDataFrame cols)-freeze df = case validateSchema @cols df of- Left _ -> Nothing- Right _ -> Just (TDF df)---- | Like 'freeze' but returns a descriptive error message on failure.-freezeWithError ::- forall cols.- (KnownSchema cols) =>- D.DataFrame -> Either T.Text (TypedDataFrame cols)-freezeWithError df = case validateSchema @cols df of- Left err -> Left err- Right _ -> Right (TDF df)--{- | Unwrap a typed DataFrame back to the untyped representation.-Always safe; discards type information.--}-thaw :: TypedDataFrame cols -> D.DataFrame-thaw (TDF df) = df--{- | Wrap an untyped DataFrame without any validation.-Used internally after delegation where the library guarantees schema correctness.--}-unsafeFreeze :: D.DataFrame -> TypedDataFrame cols-unsafeFreeze = TDF--validateSchema ::- forall cols.- (KnownSchema cols) =>- D.DataFrame -> Either T.Text ()-validateSchema df = mapM_ checkCol (schemaEvidence @cols)- where- checkCol :: (T.Text, SomeTypeRep) -> Either T.Text ()- checkCol (name, expectedRep) = case D.getColumn name df of- Nothing ->- Left $- "Column '"- <> name- <> "' not found in DataFrame. "- <> "Available columns: "- <> T.pack (show (columnNames df))- Just col ->- if matchesType expectedRep col- then Right ()- else- Left $- "Type mismatch on column '"- <> name- <> "': expected "- <> T.pack (show expectedRep)- <> ", got "- <> T.pack (C.columnTypeString col)--{- | Check if a Column's element type matches the expected SomeTypeRep.-For nullable columns (those with a bitmap), @Maybe a@ in the schema matches-a column whose inner type is @a@, since we store nullable data as-@BoxedColumn (Just bm) a@ or @UnboxedColumn (Just bm) a@ rather than-@Column (Maybe a)@.--}-matchesType :: SomeTypeRep -> C.Column -> Bool-matchesType expected col =- let expectedStr = show expected- colTypeStr = C.columnTypeString col- in expectedStr == colTypeStr- || ( C.hasMissing col -- nullable column: schema says "Maybe X", column stores "X" with a bitmap- && Just colTypeStr == stripPrefix "Maybe " expectedStr- )
− src/DataFrame/Typed/Generic.hs
@@ -1,205 +0,0 @@-{-# 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/Join.hs
@@ -1,72 +0,0 @@-{-# LANGUAGE AllowAmbiguousTypes #-}-{-# LANGUAGE DataKinds #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE RankNTypes #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}-{-# LANGUAGE TypeFamilies #-}--module DataFrame.Typed.Join (- -- * Typed joins- innerJoin,- leftJoin,- rightJoin,- fullOuterJoin,-) where--import GHC.TypeLits (Symbol)--import qualified DataFrame.Operations.Join as DJ--import DataFrame.Typed.Freeze (unsafeFreeze)-import DataFrame.Typed.Schema-import DataFrame.Typed.Types (TypedDataFrame (..))---- | Typed inner join on one or more key columns.-innerJoin ::- forall (keys :: [Symbol]) left right.- (AllKnownSymbol keys) =>- TypedDataFrame left ->- TypedDataFrame right ->- TypedDataFrame (InnerJoinSchema keys left right)-innerJoin (TDF l) (TDF r) =- unsafeFreeze (DJ.innerJoin keyNames r l)- where- keyNames = symbolVals @keys---- | Typed left join.-leftJoin ::- forall (keys :: [Symbol]) left right.- (AllKnownSymbol keys) =>- TypedDataFrame left ->- TypedDataFrame right ->- TypedDataFrame (LeftJoinSchema keys left right)-leftJoin (TDF l) (TDF r) =- unsafeFreeze (DJ.leftJoin keyNames l r)- where- keyNames = symbolVals @keys---- | Typed right join.-rightJoin ::- forall (keys :: [Symbol]) left right.- (AllKnownSymbol keys) =>- TypedDataFrame left ->- TypedDataFrame right ->- TypedDataFrame (RightJoinSchema keys left right)-rightJoin (TDF l) (TDF r) =- unsafeFreeze (DJ.rightJoin keyNames l r)- where- keyNames = symbolVals @keys---- | Typed full outer join.-fullOuterJoin ::- forall (keys :: [Symbol]) left right.- (AllKnownSymbol keys) =>- TypedDataFrame left ->- TypedDataFrame right ->- TypedDataFrame (FullOuterJoinSchema keys left right)-fullOuterJoin (TDF l) (TDF r) =- unsafeFreeze (DJ.fullOuterJoin keyNames r l)- where- keyNames = symbolVals @keys
− src/DataFrame/Typed/Lazy.hs
@@ -1,207 +0,0 @@-{-# LANGUAGE AllowAmbiguousTypes #-}-{-# LANGUAGE DataKinds #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE KindSignatures #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}--{- |-Module : DataFrame.Typed.Lazy-Copyright : (c) 2025-License : MIT-Stability : experimental--Type-safe lazy query pipelines.--This module combines the compile-time schema tracking of 'TypedDataFrame'-with the deferred execution of 'LazyDataFrame'. Queries are built as a-logical plan tree with phantom-typed schema tracking; execution is deferred-until 'run' is called.--@-{\-\# LANGUAGE DataKinds, TypeApplications, TypeOperators \#-\}-import qualified DataFrame.Typed.Lazy as TL-import DataFrame.Typed (Column)--type Schema = '[Column \"id\" Int, Column \"name\" Text, Column \"score\" Double]--main = do- let query = TL.scanCsv \@Schema \"data.csv\"- & TL.filter (TL.col \@\"score\" TL..>. TL.lit 0.5)- & TL.select \@'[\"id\", \"name\"]- df <- TL.run query -- TypedDataFrame '[Column \"id\" Int, Column \"name\" Text]- print df-@--}-module DataFrame.Typed.Lazy (- -- * Core type- TypedLazyDataFrame,-- -- * Data sources- scanCsv,- scanSeparated,- scanParquet,- fromDataFrame,- fromTypedDataFrame,-- -- * Schema-preserving operations- filter,- take,-- -- * Schema-modifying operations- derive,- select,-- -- * Aggregation- groupBy,- aggregate,-- -- * Joins- join,-- -- * Sort- sortBy,-- -- * Execution- run,-- -- * Re-exports for pipeline construction- module DataFrame.Typed.Expr,- module DataFrame.Typed.Types,- SortOrder (..),-) where--import Data.Kind (Type)-import Data.Proxy (Proxy (..))-import qualified Data.Text as T-import GHC.TypeLits (KnownSymbol, Symbol, symbolVal)-import Prelude hiding (filter, take)--import qualified DataFrame.Internal.Column as C-import qualified DataFrame.Internal.Expression as E-import DataFrame.Internal.Schema (Schema)-import DataFrame.Lazy.Internal.DataFrame (LazyDataFrame)-import qualified DataFrame.Lazy.Internal.DataFrame as L-import DataFrame.Lazy.Internal.LogicalPlan (SortOrder (..))-import DataFrame.Operations.Join (JoinType)-import DataFrame.Typed.Expr-import DataFrame.Typed.Freeze (unsafeFreeze)-import DataFrame.Typed.Schema-import DataFrame.Typed.Types---- | A lazy query with compile-time schema tracking.-newtype TypedLazyDataFrame (cols :: [Type]) = TLD {_unTLD :: LazyDataFrame}--instance Show (TypedLazyDataFrame cols) where- show (TLD ldf) = "TypedLazyDataFrame { " ++ show ldf ++ " }"---- | Scan a CSV file with a given schema.-scanCsv ::- Schema ->- T.Text ->- TypedLazyDataFrame cols-scanCsv schema path = TLD (L.scanCsv schema path)---- | Scan a character-separated file with a given schema.-scanSeparated ::- Char ->- Schema ->- T.Text ->- TypedLazyDataFrame cols-scanSeparated sep schema path = TLD (L.scanSeparated sep schema path)---- | Scan a Parquet file, directory, or glob pattern with a given schema.-scanParquet ::- Schema ->- T.Text ->- TypedLazyDataFrame cols-scanParquet schema path = TLD (L.scanParquet schema path)---- | Lift an already-loaded eager 'TypedDataFrame' into a lazy plan.-fromDataFrame :: TypedDataFrame cols -> TypedLazyDataFrame cols-fromDataFrame (TDF df) = TLD (L.fromDataFrame df)---- | Synonym for 'fromDataFrame'.-fromTypedDataFrame :: TypedDataFrame cols -> TypedLazyDataFrame cols-fromTypedDataFrame = fromDataFrame---- | Keep rows that satisfy the predicate.-filter :: TExpr cols Bool -> TypedLazyDataFrame cols -> TypedLazyDataFrame cols-filter (TExpr expr) (TLD ldf) = TLD (L.filter expr ldf)---- | Retain at most @n@ rows.-take :: Int -> TypedLazyDataFrame cols -> TypedLazyDataFrame cols-take n (TLD ldf) = TLD (L.take n ldf)---- | Add a computed column.-derive ::- forall name a cols.- (KnownSymbol name, C.Columnable a, AssertAbsent name cols) =>- TExpr cols a ->- TypedLazyDataFrame cols ->- TypedLazyDataFrame (Snoc cols (Column name a))-derive (TExpr expr) (TLD ldf) =- TLD (L.derive (T.pack (symbolVal (Proxy @name))) expr ldf)---- | Retain only the listed columns.-select ::- forall (names :: [Symbol]) cols.- (AllKnownSymbol names, AssertAllPresent names cols) =>- TypedLazyDataFrame cols ->- TypedLazyDataFrame (SubsetSchema names cols)-select (TLD ldf) = TLD (L.select (DataFrame.Typed.Schema.symbolVals @names) ldf)---- | A typed lazy grouped query.-newtype TypedLazyGrouped (keys :: [Symbol]) (cols :: [Type]) = TLG- { _unTLG :: ([T.Text], LazyDataFrame)- }---- | Group by key columns.-groupBy ::- forall (keys :: [Symbol]) cols.- (AllKnownSymbol keys, AssertAllPresent keys cols) =>- TypedLazyDataFrame cols ->- TypedLazyGrouped keys cols-groupBy (TLD ldf) = TLG (DataFrame.Typed.Schema.symbolVals @keys, ldf)---- | Aggregate a grouped lazy query.-aggregate ::- forall keys cols aggs.- TAgg keys cols aggs ->- TypedLazyGrouped keys cols ->- TypedLazyDataFrame (Append (GroupKeyColumns keys cols) (Reverse aggs))-aggregate tagg (TLG (keys, ldf)) =- TLD (L.groupBy keys (aggToNamedExprs tagg) ldf)---- | Join two lazy queries on a shared key column.-join ::- JoinType ->- T.Text ->- T.Text ->- TypedLazyDataFrame left ->- TypedLazyDataFrame right ->- TypedLazyDataFrame left -- TODO: compute join result schema-join jt leftKey rightKey (TLD left) (TLD right) =- TLD (L.join jt leftKey rightKey left right)---- | Sort the result by column name and direction.-sortBy ::- [(T.Text, SortOrder)] ->- TypedLazyDataFrame cols ->- TypedLazyDataFrame cols-sortBy cols (TLD ldf) = TLD (L.sortBy cols ldf)---- | Execute the lazy query and return a typed DataFrame.-run ::- forall cols.- (KnownSchema cols) =>- TypedLazyDataFrame cols ->- IO (TypedDataFrame cols)-run (TLD ldf) = unsafeFreeze <$> L.runDataFrame ldf---- | Convert TAgg to untyped named expressions for the lazy groupBy.-aggToNamedExprs :: TAgg keys cols aggs -> [(T.Text, E.UExpr)]-aggToNamedExprs TAggNil = []-aggToNamedExprs (TAggCons name (TExpr expr) rest) =- (name, E.UExpr expr) : aggToNamedExprs rest
− src/DataFrame/Typed/Operations.hs
@@ -1,378 +0,0 @@-{-# LANGUAGE AllowAmbiguousTypes #-}-{-# LANGUAGE DataKinds #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE FlexibleInstances #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE RankNTypes #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}-{-# LANGUAGE TypeFamilies #-}-{-# LANGUAGE TypeOperators #-}--module DataFrame.Typed.Operations (- -- * Schema-preserving operations- filterWhere,- filter,- filterBy,- filterAllJust,- filterJust,- filterNothing,- sortBy,- take,- takeLast,- drop,- dropLast,- range,- cube,- distinct,- sample,- shuffle,-- -- * Schema-modifying operations- derive,- impute,- select,- exclude,- rename,- renameMany,- insert,- insertColumn,- insertVector,- cloneColumn,- dropColumn,- replaceColumn,-- -- * Metadata- dimensions,- nRows,- nColumns,- columnNames,-- -- * Vertical merge- append,-) where--import Data.Proxy (Proxy (..))-import qualified Data.Text as T-import qualified Data.Vector as V-import GHC.TypeLits (KnownSymbol, Symbol, symbolVal)-import System.Random (RandomGen)-import Prelude hiding (drop, filter, take)--import qualified DataFrame.Functions as DF-import DataFrame.Internal.Column (Columnable)-import qualified DataFrame.Internal.Column as C-import qualified DataFrame.Operations.Aggregation as DA-import qualified DataFrame.Operations.Core as D-import DataFrame.Operations.Merge ()-import qualified DataFrame.Operations.Permutation as D-import qualified DataFrame.Operations.Subset as D-import qualified DataFrame.Operations.Transformations as D--import DataFrame.Typed.Freeze (unsafeFreeze)-import DataFrame.Typed.Schema-import DataFrame.Typed.Types (TExpr (..), TSortOrder (..), TypedDataFrame (..))-import qualified DataFrame.Typed.Types as T------------------------------------------------------------------------------------ Schema-preserving operations----------------------------------------------------------------------------------{- | Filter rows where a boolean expression evaluates to True.-The expression is validated against the schema at compile time.--}-filterWhere :: TExpr cols Bool -> TypedDataFrame cols -> TypedDataFrame cols-filterWhere (TExpr expr) (TDF df) = TDF (D.filterWhere expr df)---- | Filter rows by applying a predicate to a typed expression.-filter ::- (Columnable a) =>- TExpr cols a -> (a -> Bool) -> TypedDataFrame cols -> TypedDataFrame cols-filter (TExpr expr) pred' (TDF df) = TDF (D.filter expr pred' df)---- | Filter rows by a predicate on a column expression (flipped argument order).-filterBy ::- (Columnable a) =>- (a -> Bool) -> TExpr cols a -> TypedDataFrame cols -> TypedDataFrame cols-filterBy pred' (TExpr expr) (TDF df) = TDF (D.filterBy pred' expr df)--{- | Keep only rows where ALL Optional columns have Just values.-Strips 'Maybe' from all column types in the result schema.--@-df :: TDF '[Column \"x\" (Maybe Double), Column \"y\" Int]-filterAllJust df :: TDF '[Column \"x\" Double, Column \"y\" Int]-@--}-filterAllJust :: TypedDataFrame cols -> TypedDataFrame (StripAllMaybe cols)-filterAllJust (TDF df) = unsafeFreeze (D.filterAllJust df)--{- | Keep only rows where the named column has Just values.-Strips 'Maybe' from that column's type in the result schema.--@-filterJust \@\"x\" df-@--}-filterJust ::- forall name cols.- ( KnownSymbol name- , AssertPresent name cols- ) =>- TypedDataFrame cols -> TypedDataFrame (StripMaybeAt name cols)-filterJust (TDF df) = unsafeFreeze (D.filterJust colName df)- where- colName = T.pack (symbolVal (Proxy @name))--{- | Keep only rows where the named column has Nothing.-Schema is preserved (column types unchanged, just fewer rows).--}-filterNothing ::- forall name cols.- ( KnownSymbol name- , AssertPresent name cols- ) =>- TypedDataFrame cols -> TypedDataFrame cols-filterNothing (TDF df) = TDF (D.filterNothing colName df)- where- colName = T.pack (symbolVal (Proxy @name))--{- | Sort by the given typed sort orders.-Sort orders reference columns that are validated against the schema.--}-sortBy :: [TSortOrder cols] -> TypedDataFrame cols -> TypedDataFrame cols-sortBy ords (TDF df) = TDF (D.sortBy (map toUntypedSort ords) df)- where- toUntypedSort :: TSortOrder cols -> D.SortOrder- toUntypedSort (Asc (TExpr e)) = D.Asc e- toUntypedSort (Desc (TExpr e)) = D.Desc e---- | Take the first @n@ rows.-take :: Int -> TypedDataFrame cols -> TypedDataFrame cols-take n (TDF df) = TDF (D.take n df)---- | Take the last @n@ rows.-takeLast :: Int -> TypedDataFrame cols -> TypedDataFrame cols-takeLast n (TDF df) = TDF (D.takeLast n df)---- | Drop the first @n@ rows.-drop :: Int -> TypedDataFrame cols -> TypedDataFrame cols-drop n (TDF df) = TDF (D.drop n df)---- | Drop the last @n@ rows.-dropLast :: Int -> TypedDataFrame cols -> TypedDataFrame cols-dropLast n (TDF df) = TDF (D.dropLast n df)---- | Take rows in the given range (start, end).-range :: (Int, Int) -> TypedDataFrame cols -> TypedDataFrame cols-range r (TDF df) = TDF (D.range r df)---- | Take a sub-cube of the DataFrame.-cube :: (Int, Int) -> TypedDataFrame cols -> TypedDataFrame cols-cube c (TDF df) = TDF (D.cube c df)---- | Remove duplicate rows.-distinct :: TypedDataFrame cols -> TypedDataFrame cols-distinct (TDF df) = TDF (DA.distinct df)---- | Randomly sample a fraction of rows.-sample ::- (RandomGen g) => g -> Double -> TypedDataFrame cols -> TypedDataFrame cols-sample g frac (TDF df) = TDF (D.sample g frac df)---- | Shuffle all rows randomly.-shuffle :: (RandomGen g) => g -> TypedDataFrame cols -> TypedDataFrame cols-shuffle g (TDF df) = TDF (D.shuffle g df)------------------------------------------------------------------------------------ Schema-modifying operations----------------------------------------------------------------------------------{- | Derive a new column from a typed expression. The column name must NOT-already exist in the schema (enforced at compile time via 'AssertAbsent').-The expression is validated against the current schema.--@-df' = derive \@\"total\" (col \@\"price\" * col \@\"qty\") df--- df' :: TDF (Column \"total\" Double ': originalCols)-@--}-derive ::- forall name a cols.- ( KnownSymbol name- , Columnable a- , AssertAbsent name cols- ) =>- TExpr cols a ->- TypedDataFrame cols ->- TypedDataFrame (Snoc cols (T.Column name a))-derive (TExpr expr) (TDF df) = unsafeFreeze (D.derive colName expr df)- where- colName = T.pack (symbolVal (Proxy @name))--impute ::- forall name a cols.- ( KnownSymbol name- , Columnable a- , Maybe a ~ Lookup name cols- ) =>- a ->- TypedDataFrame cols ->- TypedDataFrame (Impute name cols)-impute value (TDF df) =- unsafeFreeze- (D.derive colName (DF.fromMaybe value (DF.col @(Maybe a) colName)) df)- where- colName = T.pack (symbolVal (Proxy @name))---- | Select a subset of columns by name.-select ::- forall (names :: [Symbol]) cols.- (AllKnownSymbol names, AssertAllPresent names cols) =>- TypedDataFrame cols -> TypedDataFrame (SubsetSchema names cols)-select (TDF df) = unsafeFreeze (D.select (symbolVals @names) df)---- | Exclude columns by name.-exclude ::- forall (names :: [Symbol]) cols.- (AllKnownSymbol names) =>- TypedDataFrame cols -> TypedDataFrame (ExcludeSchema names cols)-exclude (TDF df) = unsafeFreeze (D.exclude (symbolVals @names) df)---- | Rename a column.-rename ::- forall old new cols.- (KnownSymbol old, KnownSymbol new) =>- TypedDataFrame cols -> TypedDataFrame (RenameInSchema old new cols)-rename (TDF df) = unsafeFreeze (D.rename oldName newName df)- where- oldName = T.pack (symbolVal (Proxy @old))- newName = T.pack (symbolVal (Proxy @new))---- | Rename multiple columns from a type-level list of pairs.-renameMany ::- forall (pairs :: [(Symbol, Symbol)]) cols.- (AllKnownPairs pairs) =>- TypedDataFrame cols -> TypedDataFrame (RenameManyInSchema pairs cols)-renameMany (TDF df) = unsafeFreeze (foldRenames (pairVals @pairs) df)- where- foldRenames [] df' = df'- foldRenames ((old, new) : rest) df' = foldRenames rest (D.rename old new df')---- | Insert a new column from a Foldable container.-insert ::- forall name a cols t.- ( KnownSymbol name- , Columnable a- , Foldable t- , AssertAbsent name cols- ) =>- t a -> TypedDataFrame cols -> TypedDataFrame (T.Column name a ': cols)-insert xs (TDF df) = unsafeFreeze (D.insert colName xs df)- where- colName = T.pack (symbolVal (Proxy @name))---- | Insert a raw 'Column' value.-insertColumn ::- forall name a cols.- ( KnownSymbol name- , Columnable a- , AssertAbsent name cols- ) =>- C.Column -> TypedDataFrame cols -> TypedDataFrame (T.Column name a ': cols)-insertColumn col (TDF df) = unsafeFreeze (D.insertColumn colName col df)- where- colName = T.pack (symbolVal (Proxy @name))---- | Insert a boxed 'Vector'.-insertVector ::- forall name a cols.- ( KnownSymbol name- , Columnable a- , AssertAbsent name cols- ) =>- V.Vector a -> TypedDataFrame cols -> TypedDataFrame (T.Column name a ': cols)-insertVector vec (TDF df) = unsafeFreeze (D.insertVector colName vec df)- where- colName = T.pack (symbolVal (Proxy @name))---- | Clone an existing column under a new name.-cloneColumn ::- forall old new cols.- ( KnownSymbol old- , KnownSymbol new- , AssertPresent old cols- , AssertAbsent new cols- ) =>- TypedDataFrame cols -> TypedDataFrame (T.Column new (Lookup old cols) ': cols)-cloneColumn (TDF df) = unsafeFreeze (D.cloneColumn oldName newName df)- where- oldName = T.pack (symbolVal (Proxy @old))- newName = T.pack (symbolVal (Proxy @new))---- | Drop a column by name.-dropColumn ::- forall name cols.- ( KnownSymbol name- , AssertPresent name cols- ) =>- TypedDataFrame cols -> TypedDataFrame (RemoveColumn name cols)-dropColumn (TDF df) = unsafeFreeze (D.exclude [colName] df)- where- colName = T.pack (symbolVal (Proxy @name))--{- | Replace an existing column with new values derived from a typed expression.-The column must already exist and the new type must match.--}-replaceColumn ::- forall name a cols.- ( KnownSymbol name- , Columnable a- , a ~ SafeLookup name cols- , AssertPresent name cols- ) =>- TExpr cols a -> TypedDataFrame cols -> TypedDataFrame cols-replaceColumn (TExpr expr) (TDF df) = unsafeFreeze (D.derive colName expr df)- where- colName = T.pack (symbolVal (Proxy @name))---- | Vertically merge two DataFrames with the same schema.-append :: TypedDataFrame cols -> TypedDataFrame cols -> TypedDataFrame cols-append (TDF a) (TDF b) = TDF (a <> b)------------------------------------------------------------------------------------ Metadata (pass-through)----------------------------------------------------------------------------------dimensions :: TypedDataFrame cols -> (Int, Int)-dimensions (TDF df) = D.dimensions df--nRows :: TypedDataFrame cols -> Int-nRows (TDF df) = D.nRows df--nColumns :: TypedDataFrame cols -> Int-nColumns (TDF df) = D.nColumns df--columnNames :: TypedDataFrame cols -> [T.Text]-columnNames (TDF df) = D.columnNames df------------------------------------------------------------------------------------ Internal helpers------------------------------------------------------------------------------------ | Helper class for extracting [(Text, Text)] from type-level pairs.-class AllKnownPairs (pairs :: [(Symbol, Symbol)]) where- pairVals :: [(T.Text, T.Text)]--instance AllKnownPairs '[] where- pairVals = []--instance- (KnownSymbol a, KnownSymbol b, AllKnownPairs rest) =>- AllKnownPairs ('(a, b) ': rest)- where- pairVals =- ( T.pack (symbolVal (Proxy @a))- , T.pack (symbolVal (Proxy @b))- )- : pairVals @rest
− src/DataFrame/Typed/Record.hs
@@ -1,105 +0,0 @@-{-# 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/Schema.hs
@@ -1,441 +0,0 @@-{-# LANGUAGE AllowAmbiguousTypes #-}-{-# LANGUAGE ConstraintKinds #-}-{-# LANGUAGE DataKinds #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE FlexibleInstances #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE MultiParamTypeClasses #-}-{-# LANGUAGE PolyKinds #-}-{-# LANGUAGE RankNTypes #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}-{-# LANGUAGE TypeFamilies #-}-{-# LANGUAGE TypeOperators #-}-{-# LANGUAGE UndecidableInstances #-}--module DataFrame.Typed.Schema (- -- * Type families for schema manipulation- Lookup,- SafeLookup,- HasName,- RemoveColumn,- Impute,- SubsetSchema,- ExcludeSchema,- RenameInSchema,- RenameManyInSchema,- Append,- Snoc,- Reverse,- ColumnNames,- AssertAbsent,- AssertPresent,- AssertAllPresent,- IsElem,-- -- * Maybe-stripping families- StripAllMaybe,- StripMaybeAt,-- -- * Join schema families- SharedNames,- UniqueLeft,- InnerJoinSchema,- LeftJoinSchema,- RightJoinSchema,- FullOuterJoinSchema,- WrapMaybe,- WrapMaybeColumns,- CollidingColumns,-- -- * GroupBy helpers- GroupKeyColumns,-- -- * KnownSchema class- KnownSchema (..),-- -- * Helpers- AllKnownSymbol (..),-) where--import Data.Kind (Constraint, Type)-import Data.Proxy (Proxy (..))-import qualified Data.Text as T-import Data.These (These)-import GHC.TypeLits-import Type.Reflection (SomeTypeRep, Typeable, someTypeRep)--import DataFrame.Internal.Column (Columnable)-import DataFrame.Typed.Types (Column)---- | Look up the element type of a column by name.-type family Lookup (name :: Symbol) (cols :: [Type]) :: Type where- Lookup name (Column name a ': _) = a- Lookup name (Column _ _ ': rest) = Lookup name rest- Lookup name '[] =- TypeError- ('Text "Column '" ':<>: 'Text name ':<>: 'Text "' not found in schema")--{- | Like 'Lookup', but returns a harmless fallback ('Int') instead of-'TypeError' when the column is not found. Use together with-'AssertPresent' so the error fires exactly once.--}-type family SafeLookup (name :: Symbol) (cols :: [Type]) :: Type where- SafeLookup name (Column name a ': _) = a- SafeLookup name (Column _ _ ': rest) = SafeLookup name rest- SafeLookup name '[] = Int---- | Unwrap a Maybe from a type after we impute values.-type family Impute (name :: Symbol) (cols :: [Type]) :: [Type] where- Impute name (Column name (Maybe a) ': rest) = Column name a ': rest- Impute name (Column name _ ': rest) =- TypeError- ('Text "Column '" ':<>: 'Text name ':<>: 'Text "' is not of kind Maybe *")- Impute name (col ': rest) = col ': Impute name rest- Impute name '[] = '[]---- | Add type to the end of a list.-type family Snoc (xs :: [k]) (x :: k) :: [k] where- Snoc '[] x = '[x]- Snoc (y ': ys) x = y ': Snoc ys x---- | Check whether a column name exists in a schema (type-level Bool).-type family HasName (name :: Symbol) (cols :: [Type]) :: Bool where- HasName name (Column name _ ': _) = 'True- HasName name (Column _ _ ': rest) = HasName name rest- HasName name '[] = 'False---- | Remove a column by name from a schema.-type family RemoveColumn (name :: Symbol) (cols :: [Type]) :: [Type] where- RemoveColumn name (Column name _ ': rest) = rest- RemoveColumn name (col ': rest) = col ': RemoveColumn name rest- RemoveColumn name '[] = '[]---- | Select a subset of columns by a list of names.-type family SubsetSchema (names :: [Symbol]) (cols :: [Type]) :: [Type] where- SubsetSchema '[] cols = '[]- SubsetSchema (n ': ns) cols = Column n (Lookup n cols) ': SubsetSchema ns cols---- | Exclude columns by a list of names.-type family ExcludeSchema (names :: [Symbol]) (cols :: [Type]) :: [Type] where- ExcludeSchema names '[] = '[]- ExcludeSchema names (Column n a ': rest) =- ExcludeSchemaHelper (IsElem n names) n a names rest--type family- ExcludeSchemaHelper- (found :: Bool)- (n :: Symbol)- (a :: Type)- (names :: [Symbol])- (rest :: [Type]) ::- [Type]- where- ExcludeSchemaHelper 'True n a names rest = ExcludeSchema names rest- ExcludeSchemaHelper 'False n a names rest =- Column n a ': ExcludeSchema names rest---- | Type-level elem for Symbols-type family IsElem (x :: Symbol) (xs :: [Symbol]) :: Bool where- IsElem x '[] = 'False- IsElem x (x ': _) = 'True- IsElem x (_ ': xs) = IsElem x xs---- | Rename a column in the schema.-type family RenameInSchema (old :: Symbol) (new :: Symbol) (cols :: [Type]) :: [Type] where- RenameInSchema old new (Column old a ': rest) = Column new a ': rest- RenameInSchema old new (col ': rest) = col ': RenameInSchema old new rest- RenameInSchema old new '[] =- TypeError- ('Text "Cannot rename: column '" ':<>: 'Text old ':<>: 'Text "' not found")---- | Rename multiple columns.-type family RenameManyInSchema (pairs :: [(Symbol, Symbol)]) (cols :: [Type]) :: [Type] where- RenameManyInSchema '[] cols = cols- RenameManyInSchema ('(old, new) ': rest) cols =- RenameManyInSchema rest (RenameInSchema old new cols)---- | Append two type-level lists.-type family Append (xs :: [k]) (ys :: [k]) :: [k] where- Append '[] ys = ys- Append (x ': xs) ys = x ': Append xs ys---- | Reverse a type-level list.-type family Reverse (xs :: [Type]) :: [Type] where- Reverse xs = ReverseAcc xs '[]--type family ReverseAcc (xs :: [Type]) (acc :: [Type]) :: [Type] where- ReverseAcc '[] acc = acc- ReverseAcc (x ': xs) acc = ReverseAcc xs (x ': acc)---- | Extract column names as a type-level list of Symbols.-type family ColumnNames (cols :: [Type]) :: [Symbol] where- ColumnNames '[] = '[]- ColumnNames (Column n _ ': rest) = n ': ColumnNames rest---- | Assert that a column name is absent from the schema (for derive/insert).-type family AssertAbsent (name :: Symbol) (cols :: [Type]) :: Constraint where- AssertAbsent name cols = AssertAbsentHelper name (HasName name cols) cols--type family- AssertAbsentHelper (name :: Symbol) (found :: Bool) (cols :: [Type]) ::- Constraint- where- AssertAbsentHelper name 'False cols = ()- AssertAbsentHelper name 'True cols =- TypeError- ( 'Text "Column '"- ':<>: 'Text name- ':<>: 'Text "' already exists in schema. "- ':<>: 'Text "Use replaceColumn to overwrite."- )---- | Assert that a column name is present in the schema.-type family AssertPresent (name :: Symbol) (cols :: [Type]) :: Constraint where- AssertPresent name cols = AssertPresentHelper name (HasName name cols) cols--type family- AssertPresentHelper (name :: Symbol) (found :: Bool) (cols :: [Type]) ::- Constraint- where- AssertPresentHelper name 'True cols = ()- AssertPresentHelper name 'False cols =- TypeError- ('Text "Column '" ':<>: 'Text name ':<>: 'Text "' not found in schema")---- | Assert that a column name is present in the schema.-type family AssertAllPresent (name :: [Symbol]) (cols :: [Type]) :: Constraint where- AssertAllPresent (name ': rest) cols =- AssertAllPresentHelper (HasName name cols) name rest cols- AssertAllPresent '[] cols = ()--type family- AssertAllPresentHelper- (found :: Bool)- (name :: Symbol)- (rest :: [Symbol])- (cols :: [Type]) ::- Constraint- where- AssertAllPresentHelper 'True name rest cols = AssertAllPresent rest cols- AssertAllPresentHelper 'False name rest cols =- TypeError- ('Text "Column '" ':<>: 'Text name ':<>: 'Text "' not found in schema")--{- | Strip 'Maybe' from all columns. Used by 'filterAllJust'.--@Column "x" (Maybe Double)@ becomes @Column "x" Double@.-@Column "y" Int@ stays @Column "y" Int@.--}-type family StripAllMaybe (cols :: [Type]) :: [Type] where- StripAllMaybe '[] = '[]- StripAllMaybe (Column n (Maybe a) ': rest) = Column n a ': StripAllMaybe rest- StripAllMaybe (Column n a ': rest) = Column n a ': StripAllMaybe rest--{- | Strip 'Maybe' from a single named column. Used by 'filterJust'.--@StripMaybeAt "x" '[Column "x" (Maybe Double), Column "y" Int]@- = @'[Column "x" Double, Column "y" Int]@--}-type family StripMaybeAt (name :: Symbol) (cols :: [Type]) :: [Type] where- StripMaybeAt name (Column name (Maybe a) ': rest) = Column name a ': rest- StripMaybeAt name (Column name a ': rest) = Column name a ': rest- StripMaybeAt name (col ': rest) = col ': StripMaybeAt name rest- StripMaybeAt name '[] =- TypeError- ('Text "Column '" ':<>: 'Text name ':<>: 'Text "' not found in schema")---- | Extract column names that appear in both schemas.-type family SharedNames (left :: [Type]) (right :: [Type]) :: [Symbol] where- SharedNames '[] right = '[]- SharedNames (Column n _ ': rest) right =- SharedNamesHelper (HasName n right) n rest right--type family- SharedNamesHelper- (found :: Bool)- (n :: Symbol)- (rest :: [Type])- (right :: [Type]) ::- [Symbol]- where- SharedNamesHelper 'True n rest right = n ': SharedNames rest right- SharedNamesHelper 'False n rest right = SharedNames rest right---- | Columns from @left@ whose names do NOT appear in @right@.-type family UniqueLeft (left :: [Type]) (rightNames :: [Symbol]) :: [Type] where- UniqueLeft '[] _ = '[]- UniqueLeft (Column n a ': rest) rn =- UniqueLeftHelper (IsElem n rn) n a rest rn--type family- UniqueLeftHelper- (found :: Bool)- (n :: Symbol)- (a :: Type)- (rest :: [Type])- (rn :: [Symbol]) ::- [Type]- where- UniqueLeftHelper 'True n a rest rn = UniqueLeft rest rn- UniqueLeftHelper 'False n a rest rn = Column n a ': UniqueLeft rest rn---- | Wrap column types in Maybe.-type family WrapMaybe (cols :: [Type]) :: [Type] where- WrapMaybe '[] = '[]- WrapMaybe (Column n a ': rest) = Column n (Maybe a) ': WrapMaybe rest---- | Wrap selected columns in Maybe by name list.-type family WrapMaybeColumns (names :: [Symbol]) (cols :: [Type]) :: [Type] where- WrapMaybeColumns names '[] = '[]- WrapMaybeColumns names (Column n a ': rest) =- WrapMaybeColumnsHelper (IsElem n names) n a names rest--type family- WrapMaybeColumnsHelper- (found :: Bool)- (n :: Symbol)- (a :: Type)- (names :: [Symbol])- (rest :: [Type]) ::- [Type]- where- WrapMaybeColumnsHelper 'True n a names rest =- Column n (Maybe a) ': WrapMaybeColumns names rest- WrapMaybeColumnsHelper 'False n a names rest =- Column n a ': WrapMaybeColumns names rest---- | Columns in left whose names collide with right (excluding keys).-type family CollidingColumns (left :: [Type]) (right :: [Type]) (keys :: [Symbol]) :: [Type] where- CollidingColumns '[] _ _ = '[]- CollidingColumns (Column n a ': rest) right keys =- CollidingColumnsHelper1 (IsElem n keys) n a rest right keys--type family- CollidingColumnsHelper1- (isKey :: Bool)- (n :: Symbol)- (a :: Type)- (rest :: [Type])- (right :: [Type])- (keys :: [Symbol]) ::- [Type]- where- CollidingColumnsHelper1 'True n a rest right keys =- CollidingColumns rest right keys- CollidingColumnsHelper1 'False n a rest right keys =- CollidingColumnsHelper2 (HasName n right) n a rest right keys--type family- CollidingColumnsHelper2- (inRight :: Bool)- (n :: Symbol)- (a :: Type)- (rest :: [Type])- (right :: [Type])- (keys :: [Symbol]) ::- [Type]- where- CollidingColumnsHelper2 'True n a rest right keys =- Column n (These a (Lookup n right)) ': CollidingColumns rest right keys- CollidingColumnsHelper2 'False n a rest right keys =- CollidingColumns rest right keys---- | Inner join result schema.-type family InnerJoinSchema (keys :: [Symbol]) (left :: [Type]) (right :: [Type]) :: [Type] where- InnerJoinSchema keys left right =- Append- (SubsetSchema keys left)- ( Append- (UniqueLeft left (Append keys (ColumnNames right)))- ( Append- (UniqueLeft right (Append keys (ColumnNames left)))- (CollidingColumns left right keys)- )- )---- | Left join result schema.-type family LeftJoinSchema (keys :: [Symbol]) (left :: [Type]) (right :: [Type]) :: [Type] where- LeftJoinSchema keys left right =- Append- (SubsetSchema keys left)- ( Append- (UniqueLeft left (Append keys (ColumnNames right)))- ( Append- (WrapMaybe (UniqueLeft right (Append keys (ColumnNames left))))- (CollidingColumns left right keys)- )- )---- | Right join result schema.-type family RightJoinSchema (keys :: [Symbol]) (left :: [Type]) (right :: [Type]) :: [Type] where- RightJoinSchema keys left right =- Append- (SubsetSchema keys right)- ( Append- (WrapMaybe (UniqueLeft left (Append keys (ColumnNames right))))- ( Append- (UniqueLeft right (Append keys (ColumnNames left)))- (CollidingColumns left right keys)- )- )---- | Full outer join result schema.-type family- FullOuterJoinSchema (keys :: [Symbol]) (left :: [Type]) (right :: [Type]) ::- [Type]- where- FullOuterJoinSchema keys left right =- Append- (WrapMaybe (SubsetSchema keys left))- ( Append- (WrapMaybe (UniqueLeft left (Append keys (ColumnNames right))))- ( Append- (WrapMaybe (UniqueLeft right (Append keys (ColumnNames left))))- (CollidingColumns left right keys)- )- )---- | Extract Column entries from a schema whose names appear in @keys@.-type family GroupKeyColumns (keys :: [Symbol]) (cols :: [Type]) :: [Type] where- GroupKeyColumns keys '[] = '[]- GroupKeyColumns keys (Column n a ': rest) =- GroupKeyColumnsHelper (IsElem n keys) n a keys rest--type family- GroupKeyColumnsHelper- (found :: Bool)- (n :: Symbol)- (a :: Type)- (keys :: [Symbol])- (rest :: [Type]) ::- [Type]- where- GroupKeyColumnsHelper 'True n a keys rest =- Column n a ': GroupKeyColumns keys rest- GroupKeyColumnsHelper 'False n a keys rest = GroupKeyColumns keys rest---- | Provides runtime evidence of a schema: a list of (name, TypeRep) pairs.-class KnownSchema (cols :: [Type]) where- schemaEvidence :: [(T.Text, SomeTypeRep)]--instance KnownSchema '[] where- schemaEvidence = []--instance- (KnownSymbol name, Typeable a, Columnable a, KnownSchema rest) =>- KnownSchema (Column name a ': rest)- where- schemaEvidence =- (T.pack (symbolVal (Proxy @name)), someTypeRep (Proxy @a))- : schemaEvidence @rest---- | A class that provides a list of 'Text' values for a type-level list of Symbols.-class AllKnownSymbol (names :: [Symbol]) where- symbolVals :: [T.Text]--instance AllKnownSymbol '[] where- symbolVals = []--instance (KnownSymbol n, AllKnownSymbol ns) => AllKnownSymbol (n ': ns) where- symbolVals = T.pack (symbolVal (Proxy @n)) : symbolVals @ns
src/DataFrame/Typed/TH.hs view
@@ -1,319 +1,27 @@-{-# LANGUAGE DataKinds #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TemplateHaskellQuotes #-}-{-# LANGUAGE TypeApplications #-}--module DataFrame.Typed.TH (- -- * Schema inference- deriveSchema,- 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'.--@--}--{- $(deriveSchema \"IrisSchema\" irisDF)--- Generates: type IrisSchema = '[Column \"sepal_length\" Double, ...]-@--}--deriveSchema :: String -> D.DataFrame -> DecsQ-deriveSchema typeName df = do- let cols = getSchemaInfo df- let names = map fst cols- case findDuplicate names of- Just dup -> fail $ "Duplicate column name in DataFrame: " ++ T.unpack dup- Nothing -> pure ()- colTypes <- mapM mkColumnType cols- let schemaType = foldr (\t acc -> PromotedConsT `AppT` t `AppT` acc) PromotedNilT colTypes- let synName = mkName typeName- pure [TySynD synName [] schemaType]--deriveSchemaFromCsvFile :: String -> String -> DecsQ-deriveSchemaFromCsvFile = deriveSchemaFromCsvFileWith D.defaultReadOptions--deriveSchemaFromCsvFileWith :: D.ReadOptions -> String -> String -> DecsQ-deriveSchemaFromCsvFileWith opts typeName path = do- df <- liftIO (D.readSeparated opts path)- deriveSchema typeName df--getSchemaInfo :: D.DataFrame -> [(T.Text, String)]-getSchemaInfo df =- let orderedNames =- map fst $- L.sortBy (\(_, a) (_, b) -> compare a b) $- M.toList (D.columnIndices df)- in map (\name -> (name, getColumnTypeStr name df)) orderedNames--getColumnTypeStr :: T.Text -> D.DataFrame -> String-getColumnTypeStr name df = case D.getColumn name df of- Just col -> C.columnTypeString col- Nothing -> error $ "Column not found: " ++ T.unpack name--mkColumnType :: (T.Text, String) -> Q Type-mkColumnType (name, tyStr) = do- ty <- parseTypeString tyStr- let nameLit = LitT (StrTyLit (T.unpack name))- pure $ ConT ''Column `AppT` nameLit `AppT` ty--parseTypeString :: String -> Q Type-parseTypeString "Int" = pure $ ConT ''Int-parseTypeString "Double" = pure $ ConT ''Double-parseTypeString "Float" = pure $ ConT ''Float-parseTypeString "Bool" = pure $ ConT ''Bool-parseTypeString "Char" = pure $ ConT ''Char-parseTypeString "String" = pure $ ConT ''String-parseTypeString "Text" = pure $ ConT ''T.Text-parseTypeString "Integer" = pure $ ConT ''Integer-parseTypeString s- | "Maybe " `L.isPrefixOf` s = do- inner <- parseTypeString (L.drop 6 s)- pure $ ConT ''Maybe `AppT` inner-parseTypeString s = fail $ "Unsupported column type in schema inference: " ++ s--findDuplicate :: (Eq a) => [a] -> Maybe a-findDuplicate [] = Nothing-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+{-# LANGUAGE CPP #-} -{- | Derive a schema type synonym and a 'HasSchema' instance from a single-record-ADT, using 'defaultSchemaOptions'.+{- |+Module : DataFrame.Typed.TH+License : MIT -@-data Order = Order- { orderId :: Int64- , region :: Text- , amount :: Double- } deriving (Show, Eq)+Backwards-compatibility re-export hub for the split @DataFrame.Typed.TH.*@+modules: -\$(deriveSchemaFromType ''Order)+ * "DataFrame.Typed.TH.Records" — record-based schema derivation (no file IO)+ * "DataFrame.Typed.TH.CSV" — CSV-file-based schema derivation (requires @-fwith-csv@) --- expands to:--- type OrderSchema =--- '[Column "order_id" Int64, Column "region" Text, Column "amount" Double]------ instance HasSchema Order OrderSchema where--- toColumns rs = ...--- fromColumns df = ...-@+These live in @dataframe-th@ and @dataframe-csv-th@ respectively. The CSV+re-export is guarded with the @CSV_TH@ CPP define that the+meta-@dataframe@ package sets based on its cabal flags. -}-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) []+module DataFrame.Typed.TH (+ module DataFrame.Typed.TH.Records,+#ifdef WITH_CSV_TH+ module DataFrame.Typed.TH.CSV,+#endif+) where -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) []+import DataFrame.Typed.TH.Records+#ifdef WITH_CSV_TH+import DataFrame.Typed.TH.CSV+#endif
− src/DataFrame/Typed/Types.hs
@@ -1,114 +0,0 @@-{-# LANGUAGE AllowAmbiguousTypes #-}-{-# LANGUAGE DataKinds #-}-{-# LANGUAGE ExistentialQuantification #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE KindSignatures #-}-{-# LANGUAGE RankNTypes #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeOperators #-}--module DataFrame.Typed.Types (- -- * Core phantom-typed wrapper- TypedDataFrame (..),-- -- * Column phantom type (no constructors)- Column,-- -- * Typed expressions (schema-validated)- TExpr (..),-- -- * Typed sort orders- TSortOrder (..),-- -- * Grouped typed dataframe- TypedGrouped (..),-- -- * Typed aggregation builder (Option B)- TAgg (..),- taggToNamedExprs,-- -- * Re-export These- These (..),-) where--import Data.Kind (Type)-import Data.These (These (..))-import GHC.TypeLits (Symbol)--import qualified Data.Text as T-import DataFrame.Internal.Column (Columnable)-import qualified DataFrame.Internal.DataFrame as D-import DataFrame.Internal.Expression (Expr, NamedExpr, UExpr (..))--{- | A phantom-typed wrapper over the untyped 'DataFrame'.--The type parameter @cols@ is a type-level list of @Column name ty@ entries-that tracks the schema at compile time. All operations delegate to the-untyped core at runtime and update the phantom type at compile time.--}-newtype TypedDataFrame (cols :: [Type]) = TDF {unTDF :: D.DataFrame}--instance Show (TypedDataFrame cols) where- show (TDF df) = show df--instance Eq (TypedDataFrame cols) where- (TDF a) == (TDF b) = a == b--{- | A phantom type that pairs a type-level column name ('Symbol')-with its element type. Has no value-level constructors — used-purely at the type level to describe schemas.--}-data Column (name :: Symbol) (a :: Type)--{- | A typed expression validated against schema @cols@, producing values of type @a@.--Unlike the untyped 'Expr a', a 'TExpr' can only be constructed through-type-safe combinators ('col', 'lit', arithmetic operations) that verify-column references exist in the schema with the correct type.--Use 'unTExpr' to extract the underlying 'Expr' for delegation to the untyped API.--}-newtype TExpr (cols :: [Type]) a = TExpr {unTExpr :: Expr a}---- | A typed sort order validated against schema @cols@.-data TSortOrder (cols :: [Type]) where- Asc :: (Columnable a, Ord a) => TExpr cols a -> TSortOrder cols- Desc :: (Columnable a, Ord a) => TExpr cols a -> TSortOrder cols---- | A phantom-typed wrapper over 'GroupedDataFrame'.-newtype TypedGrouped (keys :: [Symbol]) (cols :: [Type])- = TGD {unTGD :: D.GroupedDataFrame}--{- | 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'.--@-as \@\"total\" (F.sum salary)- . as \@\"avg_age\" (F.mean age)-@--}-data TAgg (keys :: [Symbol]) (cols :: [Type]) (aggs :: [Type]) where- TAggNil :: TAgg keys cols '[]- TAggCons ::- (Columnable a) =>- -- | column name- T.Text ->- -- | typed aggregation expression- TExpr cols a ->- -- | rest- TAgg keys cols aggs ->- TAgg keys cols (Column name a ': aggs)--{- | Extract the runtime 'NamedExpr' list from a 'TAgg', in-declaration order (reversed from the cons-built order).--}-taggToNamedExprs :: TAgg keys cols aggs -> [NamedExpr]-taggToNamedExprs = reverse . go- where- go :: TAgg keys cols aggs -> [NamedExpr]- go TAggNil = []- go (TAggCons name (TExpr expr) rest) = (name, UExpr expr) : go rest
tests/Operations/Join.hs view
@@ -7,9 +7,9 @@ import Assertions (assertExpectException) import Control.Exception (evaluate) import Data.Text (Text, unpack)-import Data.These import qualified DataFrame as D import qualified DataFrame.Functions as F+import DataFrame.Internal.Types (These (..)) import DataFrame.Operations.Join import qualified DataFrame.Typed as DT import Test.HUnit