diff --git a/dataframe.cabal b/dataframe.cabal
--- a/dataframe.cabal
+++ b/dataframe.cabal
@@ -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
diff --git a/ffi/DataFrame/IO/Arrow.hs b/ffi/DataFrame/IO/Arrow.hs
--- a/ffi/DataFrame/IO/Arrow.hs
+++ b/ffi/DataFrame/IO/Arrow.hs
@@ -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
diff --git a/src/DataFrame.hs b/src/DataFrame.hs
--- a/src/DataFrame.hs
+++ b/src/DataFrame.hs
@@ -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,
diff --git a/src/DataFrame/DecisionTree.hs b/src/DataFrame/DecisionTree.hs
deleted file mode 100644
--- a/src/DataFrame/DecisionTree.hs
+++ /dev/null
@@ -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)
diff --git a/src/DataFrame/Display.hs b/src/DataFrame/Display.hs
deleted file mode 100644
--- a/src/DataFrame/Display.hs
+++ /dev/null
@@ -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)
diff --git a/src/DataFrame/Display/Terminal/Colours.hs b/src/DataFrame/Display/Terminal/Colours.hs
deleted file mode 100644
--- a/src/DataFrame/Display/Terminal/Colours.hs
+++ /dev/null
@@ -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"
diff --git a/src/DataFrame/Display/Terminal/Plot.hs b/src/DataFrame/Display/Terminal/Plot.hs
deleted file mode 100644
--- a/src/DataFrame/Display/Terminal/Plot.hs
+++ /dev/null
@@ -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"
diff --git a/src/DataFrame/Display/Terminal/PrettyPrint.hs b/src/DataFrame/Display/Terminal/PrettyPrint.hs
deleted file mode 100644
--- a/src/DataFrame/Display/Terminal/PrettyPrint.hs
+++ /dev/null
@@ -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 <> "
" <> 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
diff --git a/src/DataFrame/Display/Web/Plot.hs b/src/DataFrame/Display/Web/Plot.hs
deleted file mode 100644
--- a/src/DataFrame/Display/Web/Plot.hs
+++ /dev/null
@@ -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
- [ "\n"
- , "\n"
- , "\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"
- 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)
diff --git a/src/DataFrame/Errors.hs b/src/DataFrame/Errors.hs
deleted file mode 100644
--- a/src/DataFrame/Errors.hs
+++ /dev/null
@@ -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)
- ]
diff --git a/src/DataFrame/Functions.hs b/src/DataFrame/Functions.hs
deleted file mode 100644
--- a/src/DataFrame/Functions.hs
+++ /dev/null
@@ -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
diff --git a/src/DataFrame/IO/CSV.hs b/src/DataFrame/IO/CSV.hs
deleted file mode 100644
--- a/src/DataFrame/IO/CSV.hs
+++ /dev/null
@@ -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
diff --git a/src/DataFrame/IO/JSON.hs b/src/DataFrame/IO/JSON.hs
deleted file mode 100644
--- a/src/DataFrame/IO/JSON.hs
+++ /dev/null
@@ -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
diff --git a/src/DataFrame/IO/Parquet.hs b/src/DataFrame/IO/Parquet.hs
deleted file mode 100644
--- a/src/DataFrame/IO/Parquet.hs
+++ /dev/null
@@ -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]
diff --git a/src/DataFrame/IO/Parquet/Binary.hs b/src/DataFrame/IO/Parquet/Binary.hs
deleted file mode 100644
--- a/src/DataFrame/IO/Parquet/Binary.hs
+++ /dev/null
@@ -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
diff --git a/src/DataFrame/IO/Parquet/Decompress.hs b/src/DataFrame/IO/Parquet/Decompress.hs
deleted file mode 100644
--- a/src/DataFrame/IO/Parquet/Decompress.hs
+++ /dev/null
@@ -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)
diff --git a/src/DataFrame/IO/Parquet/Dictionary.hs b/src/DataFrame/IO/Parquet/Dictionary.hs
deleted file mode 100644
--- a/src/DataFrame/IO/Parquet/Dictionary.hs
+++ /dev/null
@@ -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)
diff --git a/src/DataFrame/IO/Parquet/Encoding.hs b/src/DataFrame/IO/Parquet/Encoding.hs
deleted file mode 100644
--- a/src/DataFrame/IO/Parquet/Encoding.hs
+++ /dev/null
@@ -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 #-}
diff --git a/src/DataFrame/IO/Parquet/Levels.hs b/src/DataFrame/IO/Parquet/Levels.hs
deleted file mode 100644
--- a/src/DataFrame/IO/Parquet/Levels.hs
+++ /dev/null
@@ -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
diff --git a/src/DataFrame/IO/Parquet/Page.hs b/src/DataFrame/IO/Parquet/Page.hs
deleted file mode 100644
--- a/src/DataFrame/IO/Parquet/Page.hs
+++ /dev/null
@@ -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)
diff --git a/src/DataFrame/IO/Parquet/Schema.hs b/src/DataFrame/IO/Parquet/Schema.hs
deleted file mode 100644
--- a/src/DataFrame/IO/Parquet/Schema.hs
+++ /dev/null
@@ -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
diff --git a/src/DataFrame/IO/Parquet/Seeking.hs b/src/DataFrame/IO/Parquet/Seeking.hs
deleted file mode 100644
--- a/src/DataFrame/IO/Parquet/Seeking.hs
+++ /dev/null
@@ -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))
diff --git a/src/DataFrame/IO/Parquet/Thrift.hs b/src/DataFrame/IO/Parquet/Thrift.hs
deleted file mode 100644
--- a/src/DataFrame/IO/Parquet/Thrift.hs
+++ /dev/null
@@ -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
diff --git a/src/DataFrame/IO/Parquet/Time.hs b/src/DataFrame/IO/Parquet/Time.hs
deleted file mode 100644
--- a/src/DataFrame/IO/Parquet/Time.hs
+++ /dev/null
@@ -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
diff --git a/src/DataFrame/IO/Parquet/Utils.hs b/src/DataFrame/IO/Parquet/Utils.hs
deleted file mode 100644
--- a/src/DataFrame/IO/Parquet/Utils.hs
+++ /dev/null
@@ -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)
diff --git a/src/DataFrame/IO/Utils/RandomAccess.hs b/src/DataFrame/IO/Utils/RandomAccess.hs
deleted file mode 100644
--- a/src/DataFrame/IO/Utils/RandomAccess.hs
+++ /dev/null
@@ -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
diff --git a/src/DataFrame/Internal/Binary.hs b/src/DataFrame/Internal/Binary.hs
deleted file mode 100644
--- a/src/DataFrame/Internal/Binary.hs
+++ /dev/null
@@ -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 #-}
diff --git a/src/DataFrame/Internal/Column.hs b/src/DataFrame/Internal/Column.hs
deleted file mode 100644
--- a/src/DataFrame/Internal/Column.hs
+++ /dev/null
@@ -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))
diff --git a/src/DataFrame/Internal/DataFrame.hs b/src/DataFrame/Internal/DataFrame.hs
deleted file mode 100644
--- a/src/DataFrame/Internal/DataFrame.hs
+++ /dev/null
@@ -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 "
diff --git a/src/DataFrame/Internal/Expression.hs b/src/DataFrame/Internal/Expression.hs
deleted file mode 100644
--- a/src/DataFrame/Internal/Expression.hs
+++ /dev/null
@@ -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) ++ ")"
diff --git a/src/DataFrame/Internal/Grouping.hs b/src/DataFrame/Internal/Grouping.hs
deleted file mode 100644
--- a/src/DataFrame/Internal/Grouping.hs
+++ /dev/null
@@ -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)
diff --git a/src/DataFrame/Internal/Interpreter.hs b/src/DataFrame/Internal/Interpreter.hs
deleted file mode 100644
--- a/src/DataFrame/Internal/Interpreter.hs
+++ /dev/null
@@ -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
diff --git a/src/DataFrame/Internal/Nullable.hs b/src/DataFrame/Internal/Nullable.hs
deleted file mode 100644
--- a/src/DataFrame/Internal/Nullable.hs
+++ /dev/null
@@ -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))
diff --git a/src/DataFrame/Internal/Parsing.hs b/src/DataFrame/Internal/Parsing.hs
deleted file mode 100644
--- a/src/DataFrame/Internal/Parsing.hs
+++ /dev/null
@@ -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)
diff --git a/src/DataFrame/Internal/Row.hs b/src/DataFrame/Internal/Row.hs
deleted file mode 100644
--- a/src/DataFrame/Internal/Row.hs
+++ /dev/null
@@ -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)
diff --git a/src/DataFrame/Internal/Schema.hs b/src/DataFrame/Internal/Schema.hs
deleted file mode 100644
--- a/src/DataFrame/Internal/Schema.hs
+++ /dev/null
@@ -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 @\Schema :: 'Schema'@ binding suitable
- for passing to 'DataFrame.IO.CSV.readCsvWithSchema' /
- 'DataFrame.IO.CSV.readCsvWithOpts'.
-* one @\\ :: '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
diff --git a/src/DataFrame/Internal/Statistics.hs b/src/DataFrame/Internal/Statistics.hs
deleted file mode 100644
--- a/src/DataFrame/Internal/Statistics.hs
+++ /dev/null
@@ -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
- #-}
diff --git a/src/DataFrame/Internal/Types.hs b/src/DataFrame/Internal/Types.hs
deleted file mode 100644
--- a/src/DataFrame/Internal/Types.hs
+++ /dev/null
@@ -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
diff --git a/src/DataFrame/Lazy.hs b/src/DataFrame/Lazy.hs
deleted file mode 100644
--- a/src/DataFrame/Lazy.hs
+++ /dev/null
@@ -1,3 +0,0 @@
-module DataFrame.Lazy (module DataFrame.Lazy.Internal.DataFrame) where
-
-import DataFrame.Lazy.Internal.DataFrame
diff --git a/src/DataFrame/Lazy/IO/Binary.hs b/src/DataFrame/Lazy/IO/Binary.hs
deleted file mode 100644
--- a/src/DataFrame/Lazy/IO/Binary.hs
+++ /dev/null
@@ -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
diff --git a/src/DataFrame/Lazy/IO/CSV.hs b/src/DataFrame/Lazy/IO/CSV.hs
deleted file mode 100644
--- a/src/DataFrame/Lazy/IO/CSV.hs
+++ /dev/null
@@ -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 #-}
diff --git a/src/DataFrame/Lazy/Internal/DataFrame.hs b/src/DataFrame/Lazy/Internal/DataFrame.hs
deleted file mode 100644
--- a/src/DataFrame/Lazy/Internal/DataFrame.hs
+++ /dev/null
@@ -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)}
diff --git a/src/DataFrame/Lazy/Internal/Executor.hs b/src/DataFrame/Lazy/Internal/Executor.hs
deleted file mode 100644
--- a/src/DataFrame/Lazy/Internal/Executor.hs
+++ /dev/null
@@ -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)
diff --git a/src/DataFrame/Lazy/Internal/LogicalPlan.hs b/src/DataFrame/Lazy/Internal/LogicalPlan.hs
deleted file mode 100644
--- a/src/DataFrame/Lazy/Internal/LogicalPlan.hs
+++ /dev/null
@@ -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 ++ " "
- 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)
diff --git a/src/DataFrame/Lazy/Internal/Optimizer.hs b/src/DataFrame/Lazy/Internal/Optimizer.hs
deleted file mode 100644
--- a/src/DataFrame/Lazy/Internal/Optimizer.hs
+++ /dev/null
@@ -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
diff --git a/src/DataFrame/Lazy/Internal/PhysicalPlan.hs b/src/DataFrame/Lazy/Internal/PhysicalPlan.hs
deleted file mode 100644
--- a/src/DataFrame/Lazy/Internal/PhysicalPlan.hs
+++ /dev/null
@@ -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)
diff --git a/src/DataFrame/Monad.hs b/src/DataFrame/Monad.hs
deleted file mode 100644
--- a/src/DataFrame/Monad.hs
+++ /dev/null
@@ -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)
diff --git a/src/DataFrame/Operations/Aggregation.hs b/src/DataFrame/Operations/Aggregation.hs
deleted file mode 100644
--- a/src/DataFrame/Operations/Aggregation.hs
+++ /dev/null
@@ -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
diff --git a/src/DataFrame/Operations/Core.hs b/src/DataFrame/Operations/Core.hs
deleted file mode 100644
--- a/src/DataFrame/Operations/Core.hs
+++ /dev/null
@@ -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 ` 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 ` 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
diff --git a/src/DataFrame/Operations/Join.hs b/src/DataFrame/Operations/Join.hs
deleted file mode 100644
--- a/src/DataFrame/Operations/Join.hs
+++ /dev/null
@@ -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
diff --git a/src/DataFrame/Operations/Merge.hs b/src/DataFrame/Operations/Merge.hs
deleted file mode 100644
--- a/src/DataFrame/Operations/Merge.hs
+++ /dev/null
@@ -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
- }
diff --git a/src/DataFrame/Operations/Permutation.hs b/src/DataFrame/Operations/Permutation.hs
deleted file mode 100644
--- a/src/DataFrame/Operations/Permutation.hs
+++ /dev/null
@@ -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
diff --git a/src/DataFrame/Operations/Statistics.hs b/src/DataFrame/Operations/Statistics.hs
deleted file mode 100644
--- a/src/DataFrame/Operations/Statistics.hs
+++ /dev/null
@@ -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)
diff --git a/src/DataFrame/Operations/Subset.hs b/src/DataFrame/Operations/Subset.hs
deleted file mode 100644
--- a/src/DataFrame/Operations/Subset.hs
+++ /dev/null
@@ -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
diff --git a/src/DataFrame/Operations/Transformations.hs b/src/DataFrame/Operations/Transformations.hs
deleted file mode 100644
--- a/src/DataFrame/Operations/Transformations.hs
+++ /dev/null
@@ -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
diff --git a/src/DataFrame/Operations/Typing.hs b/src/DataFrame/Operations/Typing.hs
deleted file mode 100644
--- a/src/DataFrame/Operations/Typing.hs
+++ /dev/null
@@ -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)
diff --git a/src/DataFrame/Operators.hs b/src/DataFrame/Operators.hs
deleted file mode 100644
--- a/src/DataFrame/Operators.hs
+++ /dev/null
@@ -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
diff --git a/src/DataFrame/Synthesis.hs b/src/DataFrame/Synthesis.hs
deleted file mode 100644
--- a/src/DataFrame/Synthesis.hs
+++ /dev/null
@@ -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
diff --git a/src/DataFrame/TH.hs b/src/DataFrame/TH.hs
--- a/src/DataFrame/TH.hs
+++ b/src/DataFrame/TH.hs
@@ -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
diff --git a/src/DataFrame/Typed.hs b/src/DataFrame/Typed.hs
--- a/src/DataFrame/Typed.hs
+++ b/src/DataFrame/Typed.hs
@@ -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 (..),
diff --git a/src/DataFrame/Typed/Access.hs b/src/DataFrame/Typed/Access.hs
deleted file mode 100644
--- a/src/DataFrame/Typed/Access.hs
+++ /dev/null
@@ -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))
diff --git a/src/DataFrame/Typed/Aggregate.hs b/src/DataFrame/Typed/Aggregate.hs
deleted file mode 100644
--- a/src/DataFrame/Typed/Aggregate.hs
+++ /dev/null
@@ -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
diff --git a/src/DataFrame/Typed/Expr.hs b/src/DataFrame/Typed/Expr.hs
deleted file mode 100644
--- a/src/DataFrame/Typed/Expr.hs
+++ /dev/null
@@ -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
diff --git a/src/DataFrame/Typed/Freeze.hs b/src/DataFrame/Typed/Freeze.hs
deleted file mode 100644
--- a/src/DataFrame/Typed/Freeze.hs
+++ /dev/null
@@ -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
- )
diff --git a/src/DataFrame/Typed/Generic.hs b/src/DataFrame/Typed/Generic.hs
deleted file mode 100644
--- a/src/DataFrame/Typed/Generic.hs
+++ /dev/null
@@ -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)
diff --git a/src/DataFrame/Typed/Join.hs b/src/DataFrame/Typed/Join.hs
deleted file mode 100644
--- a/src/DataFrame/Typed/Join.hs
+++ /dev/null
@@ -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
diff --git a/src/DataFrame/Typed/Lazy.hs b/src/DataFrame/Typed/Lazy.hs
deleted file mode 100644
--- a/src/DataFrame/Typed/Lazy.hs
+++ /dev/null
@@ -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
diff --git a/src/DataFrame/Typed/Operations.hs b/src/DataFrame/Typed/Operations.hs
deleted file mode 100644
--- a/src/DataFrame/Typed/Operations.hs
+++ /dev/null
@@ -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
diff --git a/src/DataFrame/Typed/Record.hs b/src/DataFrame/Typed/Record.hs
deleted file mode 100644
--- a/src/DataFrame/Typed/Record.hs
+++ /dev/null
@@ -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)
diff --git a/src/DataFrame/Typed/Schema.hs b/src/DataFrame/Typed/Schema.hs
deleted file mode 100644
--- a/src/DataFrame/Typed/Schema.hs
+++ /dev/null
@@ -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
diff --git a/src/DataFrame/Typed/TH.hs b/src/DataFrame/Typed/TH.hs
--- a/src/DataFrame/Typed/TH.hs
+++ b/src/DataFrame/Typed/TH.hs
@@ -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: @\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
diff --git a/src/DataFrame/Typed/Types.hs b/src/DataFrame/Typed/Types.hs
deleted file mode 100644
--- a/src/DataFrame/Typed/Types.hs
+++ /dev/null
@@ -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
diff --git a/tests/Operations/Join.hs b/tests/Operations/Join.hs
--- a/tests/Operations/Join.hs
+++ b/tests/Operations/Join.hs
@@ -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