diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,11 @@ # Revision history for dataframe +## 1.1.2.0 +* Safe read can now choose between Either and Maybe for error handling. +* Add countAll and over (window) functions. +* Add mkRandom to make random columns. +* Faster CSV parsing. + ## 1.1.1.0 ### New features * Add `DataFrame.Typed.Lazy` module — a type-safe lazy query pipeline combining compile-time schema tracking with deferred execution. diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -21,58 +21,348 @@ # DataFrame -A fast, safe, and intuitive DataFrame library. +Tabular data analysis in Haskell. Read CSV, Parquet, and JSON files, transform columns with a typed expression DSL, and optionally lock down your entire schema at the type level for compile-time safety. -## Why use this DataFrame library? +The library ships three API layers — all operating on the same underlying `DataFrame` type at runtime: -* Encourages concise, declarative, and composable data pipelines. -* Lets you opt into your preferred level of type safety: keep it lightweight for rapid exploration or lock it down completely for robust production pipelines. -* Delivers high performance thanks to Haskell’s optimizing compiler and efficient memory model. -* Designed for interactivity: expressive syntax, helpful error messages, and sensible defaults. -* Works seamlessly in both command-line and notebook environments—great for exploration and scripting alike. +- **Untyped** (`import qualified DataFrame as D`) — string-based column names, great for exploration and scripting. +- **Typed** (`import qualified DataFrame.Typed as T`) — phantom-type schema tracking with compile-time column validation. +- **Monadic API** — write your transformation as a self contained pipeline. -## Features -- Type-safe column operations with compile-time guarantees -- Familiar, approachable API designed to feel easy coming from other languages. -- Interactive REPL for data exploration and plotting. +## Why this library? -## Quick start -Browse through some examples in [binder](https://mybinder.org/v2/gh/mchav/ihaskell-dataframe/HEAD). +* Concise, declarative, composable data pipelines using the `|>` pipe operator. +* Choose your level of type safety: keep it lightweight for quick analysis, or lock it down for production pipelines. +* High performance from Haskell's optimizing compiler and an efficient columnar memory model with bitmap-backed nullability. +* Designed for interactivity: a custom REPL, IHaskell notebook support, terminal and web plotting, and helpful error messages. ## Install -See the [Quick Start](https://dataframe.readthedocs.io/en/latest/quick_start.html) guide for setup and installation instructions. -## Example +```bash +cabal update +cabal install dataframe +``` +To use as a dependency in a project: + +``` +build-depends: base >= 4, dataframe +``` + +Works with GHC 9.4 through 9.12. A custom REPL with all imports pre-loaded is available after installing: + +```bash +dataframe +``` + +## Quick Start + +Save this as `Example.hs` and run with `cabal run Example.hs`: + ```haskell -dataframe> df = D.fromNamedColumns [("product_id", D.fromList [1,1,2,2,3,3]), ("sales", D.fromList [100,120,50,20,40,30])] -dataframe> df ------------------- -product_id | sales ------------|------ - Int | Int ------------|------ -1 | 100 -1 | 120 -2 | 50 -2 | 20 -3 | 40 -3 | 30 +#!/usr/bin/env cabal +{- cabal: + build-depends: base >= 4, dataframe +-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} +import qualified DataFrame as D +import qualified DataFrame.Functions as F +import DataFrame.Operators + +main :: IO () +main = do + let sales = D.fromNamedColumns + [ ("product", D.fromList [1, 1, 2, 2, 3, 3 :: Int]) + , ("amount", D.fromList [100, 120, 50, 20, 40, 30 :: Int]) + ] + + -- Group by product and compute totals + print $ sales + |> D.groupBy ["product"] + |> D.aggregate [ F.sum (F.col @Int "amount") `as` "total" + , F.count (F.col @Int "amount") `as` "orders" + ] +``` + +``` +----------------------- +product | total | orders +--------|-------|------- + Int | Int | Int +--------|-------|------- +1 | 220 | 2 +2 | 70 | 2 +3 | 70 | 2 +``` + +Reading from files works the same way: + +```haskell +df <- D.readCsv "data.csv" +df <- D.readParquet "data.parquet" + +-- Hugging Face datasets +df <- D.readParquet "hf://datasets/scikit-learn/iris/default/train/0000.parquet" +``` + +## Interactive REPL + +The `dataframe` REPL comes with all imports pre-loaded. Here's a typical exploration session: + +```haskell +dataframe> df <- D.readCsv "./data/housing.csv" +dataframe> D.dimensions df +(20640, 10) + +dataframe> D.describeColumns df +------------------------------------------------------------------------ + Column Name | ## Non-null Values | ## Null Values | Type +--------------------|--------------------|----------------|------------- + Text | Int | Int | Text +--------------------|--------------------|----------------|------------- + total_bedrooms | 20433 | 207 | Maybe Double + ocean_proximity | 20640 | 0 | Text + median_house_value | 20640 | 0 | Double + median_income | 20640 | 0 | Double + households | 20640 | 0 | Double + population | 20640 | 0 | Double + total_rooms | 20640 | 0 | Double + housing_median_age | 20640 | 0 | Double + latitude | 20640 | 0 | Double + longitude | 20640 | 0 | Double +``` + +The `:declareColumns` macro generates typed column references from a dataframe, so you can use column names directly in expressions instead of writing `F.col @Double "median_income"` every time: + +```haskell dataframe> :declareColumns df -"product_id :: Expr Int" -"sales :: Expr Int" -dataframe> df |> D.groupBy [F.name product_id] |> D.aggregate [F.sum sales `as` "total_sales"] ------------------------- -product_id | total_sales ------------|------------ - Int | Int ------------|------------ -1 | 220 -2 | 70 -3 | 70 +"longitude :: Expr Double" +"latitude :: Expr Double" +"housing_median_age :: Expr Double" +"total_rooms :: Expr Double" +"total_bedrooms :: Expr (Maybe Double)" +"population :: Expr Double" +"households :: Expr Double" +"median_income :: Expr Double" +"median_house_value :: Expr Double" +"ocean_proximity :: Expr Text" + +dataframe> df |> D.groupBy ["ocean_proximity"] + |> D.aggregate [F.mean median_house_value `as` "avg_value"] +------------------------------------- + ocean_proximity | avg_value +-----------------|------------------- + Text | Double +-----------------|------------------- + <1H OCEAN | 240084.28546409807 + INLAND | 124805.39200122119 + ISLAND | 380440.0 + NEAR BAY | 259212.31179039303 + NEAR OCEAN | 249433.97742663656 ``` +Create new columns from existing ones: + +```haskell +dataframe> df |> D.derive "rooms_per_household" (total_rooms / households) |> D.take 3 +----------------------------------------------------------------------------------------------------------------- + longitude | latitude | housing_median_age | total_rooms | ... | ocean_proximity | rooms_per_household +-----------|----------|--------------------|-------------|-----|-----------------|-------------------- + Double | Double | Double | Double | ... | Text | Double +-----------|----------|--------------------|-------------|-----|-----------------|-------------------- + -122.23 | 37.88 | 41.0 | 880.0 | ... | NEAR BAY | 6.984126984126984 + -122.22 | 37.86 | 21.0 | 7099.0 | ... | NEAR BAY | 6.238137082601054 + -122.24 | 37.85 | 52.0 | 1467.0 | ... | NEAR BAY | 8.288135593220339 +``` + +Type mismatches are caught as compile errors — adding a `Double` column to a `Text` column won't silently produce garbage: + +```haskell +dataframe> df |> D.derive "nonsense" (latitude + ocean_proximity) + +:14:47: error: [GHC-83865] + • Couldn't match type 'Text' with 'Double' + Expected: Expr Double + Actual: Expr Text + • In the second argument of '(+)', namely 'ocean_proximity' + In the second argument of 'derive', namely + '(latitude + ocean_proximity)' +``` + +## Template Haskell + +For scripts and projects, Template Haskell can generate column bindings at compile time. + +### Generate column references from a CSV + +`declareColumnsFromCsvFile` reads your CSV at compile time and generates typed `Expr` bindings for every column: + +```haskell +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE OverloadedStrings #-} + +import qualified DataFrame as D +import qualified DataFrame.Functions as F +import DataFrame.Operators + +-- Reads housing.csv at compile time and generates: +-- latitude :: Expr Double +-- total_rooms :: Expr Double +-- ocean_proximity :: Expr Text +-- ... one binding per column +$(F.declareColumnsFromCsvFile "./data/housing.csv") + +main :: IO () +main = do + df <- D.readCsv "./data/housing.csv" + print $ df + |> D.derive "rooms_per_household" (total_rooms / households) + |> D.filterWhere (median_income .>. 5) + |> D.groupBy ["ocean_proximity"] + |> D.aggregate [F.mean median_house_value `as` "avg_value"] +``` + +Compare this to the manual version which requires spelling out every column name and type: + +```haskell +-- Without TH — every column needs its name and type spelled out +df |> D.derive "rooms_per_household" + (F.col @Double "total_rooms" / F.col @Double "households") + |> D.filterWhere (F.col @Double "median_income" .>. F.lit 5) +``` + +### Generate a schema type from a CSV + +`deriveSchemaFromCsvFile` generates a type synonym for use with the typed API — instead of manually writing out every column name and type: + +```haskell +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE DataKinds #-} + +import qualified DataFrame.Typed as T + +-- Generates: +-- type HousingSchema = '[ T.Column "longitude" Double +-- , T.Column "latitude" Double +-- , T.Column "total_rooms" Double +-- , ... +-- ] +$(T.deriveSchemaFromCsvFile "HousingSchema" "./data/housing.csv") +``` + +## Typed API + +When you want compile-time guarantees that column names exist and types match, wrap your `DataFrame` in a `TypedDataFrame`: + +```haskell +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE OverloadedStrings #-} + +import qualified DataFrame as D +import qualified DataFrame.Typed as T +import Data.Text (Text) +import DataFrame.Operators + +type EmployeeSchema = + '[ T.Column "name" Text + , T.Column "department" Text + , T.Column "salary" Double + ] + +main :: IO () +main = do + df <- D.readCsv "employees.csv" + case T.freeze @EmployeeSchema df of + Nothing -> putStrLn "Schema mismatch!" + Just tdf -> do + let result = tdf + |> T.derive @"bonus" (T.col @"salary" * T.lit 0.1) + |> T.filterWhere (T.col @"salary" .>. T.lit 50000) + |> T.select @'["name", "bonus"] + print (T.thaw result) +``` + +`T.freeze` validates the runtime `DataFrame` against your schema once at the boundary. After that, every column access is checked at compile time: + +```haskell +-- Typo in column name → compile error +tdf |> T.filterWhere (T.col @"slary" .>. T.lit 50000) +-- error: Column "slary" not found in schema + +-- Wrong type → compile error +tdf |> T.filterWhere (T.col @"name" .>. T.lit 50000) +-- error: Couldn't match type 'Text' with 'Double' +``` + +`filterAllJust` goes further — it strips `Maybe` from every column in the schema type, so downstream code can't accidentally treat cleaned columns as nullable: + +```haskell +-- Before: TypedDataFrame '[Column "score" (Maybe Double), Column "name" Text] +let cleaned = T.filterAllJust tdf +-- After: TypedDataFrame '[Column "score" Double, Column "name" Text] + +cleaned |> T.derive @"scaled" (T.col @"score" * T.lit 100) +``` + +## Features + +**I/O**: CSV, TSV, Parquet (Snappy, ZSTD, Gzip), JSON. Read Parquet from HTTP URLs and Hugging Face datasets (`hf://` URIs). Column projection and predicate pushdown for Parquet reads. + +**Operations**: filter, select, derive, groupBy, aggregate, joins (inner, left, right, full outer), sort, sample, stratified sample, distinct, k-fold splits. + +**Expressions**: typed column references (`F.col @Double "x"`), arithmetic, comparisons, logical operators, nullable-aware three-valued logic (`.==`, `.&&`), string matching (`like`, `regex`), casting, and user-defined functions via `lift`/`lift2`. + +**Statistics**: mean, median, mode, variance, standard deviation, percentiles, inter-quartile range, correlation, skewness, frequency tables, imputation. + +**Plotting**: terminal plots (histogram, scatter, line, bar, box, pie, heatmap, stacked bar, correlation matrix) and interactive HTML plots. + +**Lazy engine**: streaming query execution for files that don't fit in memory. Rule-based optimizer with filter fusion, predicate pushdown, and dead column elimination. Pull-based executor with configurable batch sizes. + +**Interop**: Arrow C Data Interface for zero-copy round-trips with Python and Polars. + +**ML**: decision trees (TAO algorithm), feature synthesis, k-fold cross-validation, stratified sampling. + +**Notebooks**: IHaskell integration with [pre-built Binder examples](https://mybinder.org/v2/gh/mchav/ihaskell-dataframe/HEAD). + +## Lazy Queries + +For files too large to fit in memory, `DataFrame.Lazy` provides a streaming query engine. Declare a schema, build a query plan with the same familiar operations, and `runDataFrame` runs it through an optimizer before streaming results batch-by-batch: + +```haskell +import qualified DataFrame.Lazy as L +import qualified DataFrame.Functions as F +import DataFrame.Operators +import DataFrame.Internal.Schema (Schema, schemaType) +import Data.Text (Text) + +mySchema :: Schema +mySchema = [ ("name", schemaType @Text) + , ("weight", schemaType @Double) + , ("height", schemaType @Double) + ] + +main :: IO () +main = do + result <- L.runDataFrame $ + L.scanCsv mySchema "large_file.csv" + |> L.filter (F.col @Double "height" .>. F.lit 1.7) + |> L.select ["name", "weight", "height"] + |> L.derive "bmi" (F.col @Double "weight" + / (F.col @Double "height" * F.col @Double "height")) + |> L.take 1000 + print result +``` + +The optimizer pushes the filter into the scan, drops unreferenced columns before reading, and stops pulling batches once 1000 rows have been collected. + ## Documentation -* 📚 User guide: https://dataframe.readthedocs.io/en/latest/ -* 📖 API reference: https://hackage.haskell.org/package/dataframe/docs/DataFrame.html + +* User guide: https://dataframe.readthedocs.io/en/latest/ +* API reference: https://hackage.haskell.org/package/dataframe/docs/DataFrame.html +* [Coming from pandas, Polars, dplyr, or Frames?](docs/coming_from_other_implementations.md) +* [Cookbook (SQL-style patterns)](docs/cookbook.md) +* [Tutorials](docs/tutorial.md) +* Discord: https://discord.gg/8u8SCWfrNC diff --git a/app/Synthesis.hs b/app/Synthesis.hs --- a/app/Synthesis.hs +++ b/app/Synthesis.hs @@ -16,17 +16,17 @@ import System.Random $( DT.deriveSchemaFromCsvFileWith - D.defaultReadOptions{D.safeRead = True} + D.defaultReadOptions{D.safeRead = D.MaybeRead} "TrainSchema" "./data/titanic/train.csv" ) $( DT.deriveSchemaFromCsvFileWith - D.defaultReadOptions{D.safeRead = True} + D.defaultReadOptions{D.safeRead = D.MaybeRead} "TestSchema" "./data/titanic/test.csv" ) --- Survived is Maybe Int (safeRead = True); prediction is Int (model output). +-- Survived is Maybe Int (safeRead = MaybeRead); prediction is Int (model output). type RawPredSchema = '[DT.Column "Survived" (Maybe Int), DT.Column "prediction" Int] @@ -118,8 +118,8 @@ -- | Extract title (e.g. "Mr", "Mrs") from a full Titanic passenger name. extractTitle :: T.Text -> T.Text -extractTitle name = - case filter (T.isSuffixOf ".") (T.words name) of +extractTitle fullName = + case filter (T.isSuffixOf ".") (T.words fullName) of (w : _) -> T.dropEnd 1 w [] -> "" @@ -133,10 +133,10 @@ DT.unsafeFreeze @RawPredSchema $ df |> D.select ["Survived", "prediction"] survived = DT.col @"Survived" - pred = DT.col @"prediction" + predCol = DT.col @"prediction" count expr = fromIntegral (DT.nRows (DT.filterWhere expr tdf)) - tp = count ((survived DT..==. DT.lit 1) DT..&&. (pred DT..==. DT.lit 1)) - tn = count ((survived DT..==. DT.lit 0) DT..&&. (pred DT..==. DT.lit 0)) - fp = count ((survived DT..==. DT.lit 0) DT..&&. (pred DT..==. DT.lit 1)) - fn = count ((survived DT..==. DT.lit 1) DT..&&. (pred DT..==. DT.lit 0)) + tp = count ((survived DT..==. DT.lit 1) DT..&&. (predCol DT..==. DT.lit 1)) + tn = count ((survived DT..==. DT.lit 0) DT..&&. (predCol DT..==. DT.lit 0)) + fp = count ((survived DT..==. DT.lit 0) DT..&&. (predCol DT..==. DT.lit 1)) + fn = count ((survived DT..==. DT.lit 1) DT..&&. (predCol DT..==. DT.lit 0)) in (tp + tn) / (tp + tn + fp + fn) diff --git a/dataframe.cabal b/dataframe.cabal --- a/dataframe.cabal +++ b/dataframe.cabal @@ -1,6 +1,6 @@ cabal-version: 2.4 name: dataframe -version: 1.1.1.0 +version: 1.1.2.0 synopsis: A fast, safe, and intuitive DataFrame library. @@ -40,13 +40,8 @@ location: https://github.com/mchav/dataframe common warnings - -- TODO add more warnings, eventually -Wall ghc-options: - -Wincomplete-patterns - -Wincomplete-uni-patterns - -Wunused-imports - -Wunused-packages - -Wunused-local-binds + -Wall library import: warnings @@ -58,6 +53,7 @@ DataFrame.Display.Web.Plot, DataFrame.Internal.Types, DataFrame.Internal.Expression, + DataFrame.Internal.Grouping, DataFrame.Internal.Interpreter, DataFrame.Internal.Nullable, DataFrame.Internal.Parsing, @@ -263,6 +259,7 @@ Operations.Nullable, Operations.Provenance, Operations.ReadCsv, + Operations.Window, Operations.WriteCsv, Operations.Shuffle, Operations.Sort, diff --git a/src/DataFrame.hs b/src/DataFrame.hs --- a/src/DataFrame.hs +++ b/src/DataFrame.hs @@ -275,6 +275,7 @@ hasElemType, hasMissing, isNumeric, + mkRandom, toList, toVector, ) @@ -388,7 +389,9 @@ ) import DataFrame.Operations.Typing as Typing ( ParseOptions (..), + SafeReadMode (..), defaultParseOptions, + effectiveSafeRead, parseDefaults, ) import DataFrame.Operators as Operators diff --git a/src/DataFrame/DecisionTree.hs b/src/DataFrame/DecisionTree.hs --- a/src/DataFrame/DecisionTree.hs +++ b/src/DataFrame/DecisionTree.hs @@ -424,7 +424,7 @@ Int -> -- Row index Tree a -> a -predictWithTree target df idx (Leaf v) = v +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 @@ -473,13 +473,13 @@ majorityValueFromIndices target df indices = case interpret @a df (Col target) of Left e -> throw e - Right (TColumn col) -> - case toVector @a col of + 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 acc) + (\acc i -> M.insertWith (+) (vals V.! i) (1 :: Int) acc) M.empty indices in if M.null counts @@ -710,13 +710,15 @@ 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))] + 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))] + STrue -> + [NMaybeDouble (F.whenPresent (realToFrac @b @Double) (Col @(Maybe b) colName))] SFalse -> [] _ -> [] @@ -775,7 +777,9 @@ 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 + case withOrdFrom @a + ords + (map (Lit . Just . (`percentileOrd'` column)) [1, 25, 75, 99]) of Just ps -> map (F.lift2 (==) (Col @(Maybe a) colName)) ps Nothing -> [] (UnboxedColumn _ (_ :: VU.Vector a)) -> [] @@ -794,7 +798,7 @@ ] where colConds (!l, !r) = case (unsafeGetColumn l df, unsafeGetColumn r df) of - ( BoxedColumn Nothing (col1 :: V.Vector a) + ( BoxedColumn Nothing (_col1 :: V.Vector a) , BoxedColumn Nothing (_ :: V.Vector b) ) -> case testEquality (typeRep @a) (typeRep @b) of @@ -826,7 +830,7 @@ 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) probs) + 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 = @@ -904,7 +908,7 @@ Right vals -> let counts = V.foldl' - (\acc i -> M.insertWith (+) (vals V.! i) 1 acc) + (\acc i -> M.insertWith (+) (vals V.! i) (1 :: Int) acc) M.empty indices total = fromIntegral (V.length indices) :: Double diff --git a/src/DataFrame/Display/Terminal/Plot.hs b/src/DataFrame/Display/Terminal/Plot.hs --- a/src/DataFrame/Display/Terminal/Plot.hs +++ b/src/DataFrame/Display/Terminal/Plot.hs @@ -171,7 +171,7 @@ numericCols ) numericCols - print (zip [0 ..] numericCols) + print (zip [(0 :: Int) ..] numericCols) T.putStrLn $ heatmap correlations (defPlot{plotTitle = "Correlation Matrix"}) where correlation xs ys = @@ -179,8 +179,8 @@ 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 | x <- xs] / n - stdY = sqrt $ sum [(y - meanY) ^ 2 | y <- 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 () @@ -254,7 +254,7 @@ M.toList $ M.fromListWith (+) - [(g <> " - " <> v, 1) | (g, v) <- pairs] + [(g <> " - " <> v, 1 :: Int) | (g, v) <- pairs] finalCounts = groupWithOther n [(k, fromIntegral v) | (k, v) <- counts] T.putStrLn $ bars finalCounts (plotSettings config) @@ -388,7 +388,7 @@ countByShow xs = map (Data.Bifunctor.bimap T.pack fromIntegral) $ M.toList $ - L.foldl' (\acc x -> M.insertWith (+) (show x) 1 acc) M.empty xs + 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 @@ -579,7 +579,7 @@ let groups = extractStringColumn groupCol df vals = extractStringColumn valCol df combined = zipWith (\g v -> g <> " - " <> v) groups vals - counts = M.toList $ M.fromListWith (+) [(c, 1) | c <- combined] + 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) diff --git a/src/DataFrame/Display/Web/Plot.hs b/src/DataFrame/Display/Web/Plot.hs --- a/src/DataFrame/Display/Web/Plot.hs +++ b/src/DataFrame/Display/Web/Plot.hs @@ -447,7 +447,7 @@ let values = extractNumericColumn colName df labels' = if length values > 20 - then take 20 ["Item " <> T.pack (show i) | i <- [1 ..]] + 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'] @@ -794,7 +794,7 @@ M.toList $ M.fromListWith (+) - [(g <> " - " <> v, 1) | (g, v) <- pairs] + [(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] @@ -946,7 +946,7 @@ countByShow xs = map (Data.Bifunctor.bimap T.pack fromIntegral) $ M.toList $ - L.foldl' (\acc x -> M.insertWith (+) (show x) 1 acc) M.empty xs + 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 = diff --git a/src/DataFrame/Functions.hs b/src/DataFrame/Functions.hs --- a/src/DataFrame/Functions.hs +++ b/src/DataFrame/Functions.hs @@ -242,6 +242,11 @@ {-# 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] #-} @@ -257,7 +262,7 @@ ( fst . L.maximumBy (compare `on` snd) . M.toList - . V.foldl' (\m e -> M.insertWith (+) e 1 m) M.empty + . V.foldl' (\m e -> M.insertWith (+) e (1 :: Int) m) M.empty ) ) {-# SPECIALIZE mode :: Expr Double -> Expr Double #-} @@ -581,6 +586,23 @@ (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 diff --git a/src/DataFrame/IO/CSV.hs b/src/DataFrame/IO/CSV.hs --- a/src/DataFrame/IO/CSV.hs +++ b/src/DataFrame/IO/CSV.hs @@ -29,6 +29,7 @@ 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 @@ -172,14 +173,49 @@ | 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 :: Bool - -- ^ Whether to partially parse values into `Maybe`/`Either`. (default: True) + , 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. @@ -198,6 +234,19 @@ -- ^ 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 @@ -219,12 +268,16 @@ ReadOptions { headerSpec = UseFirstRow , typeSpec = InferFromSample 100 - , safeRead = False + , 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. @@ -304,13 +357,23 @@ (sampleRow, _) <- peekStream rowsToProcess builderCols <- initializeColumns columnNames (V.toList sampleRow) opts let !builderColsV = V.fromList builderCols - processStream - (missingIndicators opts) - rowsToProcess - builderColsV - (numColumns opts) + 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.mapM (finalizeBuilderColumn opts) builderColsV + frozenCols <- + V.zipWithM + (\name bc -> finalizeBuilderColumn (resolveMode name) opts bc) + colNamesV + builderColsV let numRows = maybe 0 columnLength (frozenCols V.!? 0) let df = @@ -319,11 +382,11 @@ (M.fromList (zip columnNames [0 ..])) (numRows, V.length frozenCols) M.empty -- TODO give typed column references - pure $ parseWithTypes (safeRead opts) (schemaTypeMap (typeSpec opts)) df + 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) +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 @@ -332,7 +395,12 @@ 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 @@ -356,8 +424,8 @@ 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 () +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 () @@ -374,14 +442,14 @@ Nothing -> appendPagedUnboxedVector gv 0.0 >> appendPagedUnboxedVector valid 0 BuilderText gv valid -> do let !val = T.strip (TE.decodeUtf8Lenient bs') - if val `elem` missing - then appendPagedVector gv T.empty >> appendPagedUnboxedVector valid 0 - else appendPagedVector gv val >> appendPagedUnboxedVector valid 1 + appendPagedVector gv val + appendPagedUnboxedVector valid (if val `elem` missing then 0 else 1) BuilderBS gv valid -> do let !bs'' = C.strip bs' - if TE.decodeUtf8Lenient bs'' `elem` missing - then appendPagedVector gv BS.empty >> appendPagedUnboxedVector valid 0 - else appendPagedVector gv bs'' >> appendPagedUnboxedVector valid 1 + appendPagedVector gv bs'' + appendPagedUnboxedVector + valid + (if TE.decodeUtf8Lenient bs'' `elem` missing then 0 else 1) freezeBuilderColumn :: BuilderColumn -> IO Column freezeBuilderColumn (BuilderInt gv validRef) = do @@ -406,33 +474,88 @@ error "freezeBuilderColumn: BuilderBS must be finalized via finalizeBuilderColumn" -finalizeBuilderColumn :: ReadOptions -> BuilderColumn -> IO Column -finalizeBuilderColumn opts bc = do +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 opts vec valid + return $! inferColumnFromBS mode opts vec valid _ -> freezeBuilderColumn bc - return $! if safeRead opts then ensureOptional col else col + return $! case mode of + NoSafeRead -> col + MaybeRead -> ensureOptional col + EitherRead -> col inferColumnFromBS :: - ReadOptions -> V.Vector BS.ByteString -> VU.Vector Word8 -> Column -inferColumnFromBS opts vec valid = + 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 - asMaybeFull = V.generate (V.length vec) $ \i -> + -- 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 - samples = V.take sampleN asMaybeFull assumption = makeParsingAssumptionBS dfmt samples - in case assumption of - IntAssumption -> handleBSInt dfmt asMaybeFull - DoubleAssumption -> handleBSDouble asMaybeFull - BoolAssumption -> handleBSBool asMaybeFull - DateAssumption -> handleBSDate dfmt asMaybeFull - TextAssumption -> handleBSText asMaybeFull - NoAssumption -> handleBSNo dfmt asMaybeFull + 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 @@ -450,77 +573,157 @@ asMaybeDouble = V.map (>>= readByteStringDouble) asMaybe asMaybeDate = V.map (>>= readByteStringDate dfmt) asMaybe -handleBSBool :: V.Vector (Maybe BS.ByteString) -> Column -handleBSBool asMaybe - | parsableAsBool = - maybe (fromVector asMaybeBool) fromVector (sequenceA asMaybeBool) - | otherwise = handleBSText asMaybe - where - asMaybeBool = V.map (>>= readByteStringBool) asMaybe - parsableAsBool = vecSameConstructor asMaybe asMaybeBool +-- 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'. -handleBSInt :: String -> V.Vector (Maybe BS.ByteString) -> Column -handleBSInt dfmt asMaybe - | parsableAsInt = - maybe (fromVector asMaybeInt) fromVector (sequenceA asMaybeInt) - | parsableAsDouble = - maybe (fromVector asMaybeDouble) fromVector (sequenceA asMaybeDouble) - | otherwise = handleBSText asMaybe - where - asMaybeInt = V.map (>>= readByteStringInt) asMaybe - asMaybeDouble = V.map (>>= readByteStringDouble) asMaybe - parsableAsInt = - vecSameConstructor asMaybe asMaybeInt - && vecSameConstructor asMaybe asMaybeDouble - parsableAsDouble = vecSameConstructor asMaybe asMaybeDouble +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 -handleBSDouble :: V.Vector (Maybe BS.ByteString) -> Column -handleBSDouble asMaybe - | parsableAsDouble = - maybe (fromVector asMaybeDouble) fromVector (sequenceA asMaybeDouble) - | otherwise = handleBSText asMaybe - where - asMaybeDouble = V.map (>>= readByteStringDouble) asMaybe - parsableAsDouble = vecSameConstructor asMaybe asMaybeDouble +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 -handleBSDate :: String -> V.Vector (Maybe BS.ByteString) -> Column -handleBSDate dfmt asMaybe - | parsableAsDate = - maybe (fromVector asMaybeDate) fromVector (sequenceA asMaybeDate) - | otherwise = handleBSText asMaybe - where - asMaybeDate = V.map (>>= readByteStringDate dfmt) asMaybe - parsableAsDate = vecSameConstructor asMaybe asMaybeDate +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 -handleBSText :: V.Vector (Maybe BS.ByteString) -> Column -handleBSText asMaybe = - let asMaybeText = V.map (fmap TE.decodeUtf8Lenient) asMaybe - in maybe (fromVector asMaybeText) fromVector (sequenceA asMaybeText) +-- 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 -handleBSNo :: String -> V.Vector (Maybe BS.ByteString) -> Column -handleBSNo dfmt asMaybe - | V.all (== Nothing) asMaybe = - fromVector (V.map (const (Nothing :: Maybe T.Text)) asMaybe) - | parsableAsBool = - maybe (fromVector asMaybeBool) fromVector (sequenceA asMaybeBool) - | parsableAsInt = - maybe (fromVector asMaybeInt) fromVector (sequenceA asMaybeInt) - | parsableAsDouble = - maybe (fromVector asMaybeDouble) fromVector (sequenceA asMaybeDouble) - | parsableAsDate = - maybe (fromVector asMaybeDate) fromVector (sequenceA asMaybeDate) - | otherwise = handleBSText asMaybe - where - asMaybeBool = V.map (>>= readByteStringBool) asMaybe - asMaybeInt = V.map (>>= readByteStringInt) asMaybe - asMaybeDouble = V.map (>>= readByteStringDouble) asMaybe - asMaybeDate = V.map (>>= readByteStringDate dfmt) asMaybe - parsableAsBool = vecSameConstructor asMaybe asMaybeBool - parsableAsInt = - vecSameConstructor asMaybe asMaybeInt - && vecSameConstructor asMaybe asMaybeDouble - parsableAsDouble = vecSameConstructor asMaybe asMaybeDouble - parsableAsDate = vecSameConstructor asMaybe asMaybeDate +-- 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 diff --git a/src/DataFrame/IO/Parquet.hs b/src/DataFrame/IO/Parquet.hs --- a/src/DataFrame/IO/Parquet.hs +++ b/src/DataFrame/IO/Parquet.hs @@ -199,7 +199,7 @@ let schemaElements = schema fileMetadata let sNodes = parseAll (drop 1 schemaElements) let getTypeLength :: [String] -> Maybe Int32 - getTypeLength colPath = findTypeLength schemaElements colPath 0 + getTypeLength colPath = findTypeLength schemaElements colPath (0 :: Int) where findTypeLength [] _ _ = Nothing findTypeLength (s : ss) targetPath depth @@ -213,7 +213,7 @@ pathToElement _ _ _ = [] forM_ (rowGroups fileMetadata) $ \rowGroup -> do - forM_ (zip (rowGroupColumns rowGroup) [0 ..]) $ \(colChunk, colIdx) -> do + forM_ (zip (rowGroupColumns rowGroup) [(0 :: Int) ..]) $ \(colChunk, colIdx) -> do let metadata = columnMetaData colChunk let colPath = columnPathInSchema metadata let cleanPath = cleanColPath sNodes colPath @@ -434,7 +434,7 @@ Maybe Int32 -> LogicalType -> IO DI.Column -processColumnPages (maxDef, maxRep) pages pType _ maybeTypeLength lType = do +processColumnPages (maxDef, maxRep) pages pType _ maybeTypeLength _lType = do let dictPages = filter isDictionaryPage pages let dataPages = filter isDataPage pages diff --git a/src/DataFrame/IO/Parquet/Binary.hs b/src/DataFrame/IO/Parquet/Binary.hs --- a/src/DataFrame/IO/Parquet/Binary.hs +++ b/src/DataFrame/IO/Parquet/Binary.hs @@ -49,7 +49,9 @@ readIntFromBytes bs = let (n, remainder) = readVarIntFromBytes bs u = fromIntegral n :: Word32 - in (fromIntegral $ (fromIntegral (u `shiftR` 1) :: Int32) .^. (-(n .&. 1)), remainder) + in ( fromIntegral $ (fromIntegral (u `shiftR` 1) :: Int32) .^. (-(n .&. 1)) + , remainder + ) readInt32FromBytes :: BS.ByteString -> (Int32, BS.ByteString) readInt32FromBytes bs = diff --git a/src/DataFrame/IO/Parquet/Encoding.hs b/src/DataFrame/IO/Parquet/Encoding.hs --- a/src/DataFrame/IO/Parquet/Encoding.hs +++ b/src/DataFrame/IO/Parquet/Encoding.hs @@ -85,7 +85,7 @@ decodeDictIndicesV1 :: Int -> Int -> BS.ByteString -> (VU.Vector Int, BS.ByteString) -decodeDictIndicesV1 need dictCard bs = +decodeDictIndicesV1 need _dictCard bs = case BS.uncons bs of Nothing -> error "empty dictionary index stream" Just (w0, rest0) -> diff --git a/src/DataFrame/IO/Parquet/Page.hs b/src/DataFrame/IO/Parquet/Page.hs --- a/src/DataFrame/IO/Parquet/Page.hs +++ b/src/DataFrame/IO/Parquet/Page.hs @@ -1,4 +1,3 @@ -{-# LANGUAGE RecordWildCards #-} {-# LANGUAGE TypeApplications #-} module DataFrame.IO.Parquet.Page where @@ -24,13 +23,13 @@ isDataPage :: Page -> Bool isDataPage page = case pageTypeHeader (pageHeader page) of - DataPageHeader{..} -> True - DataPageHeaderV2{..} -> True + DataPageHeader{} -> True + DataPageHeaderV2{} -> True _ -> False isDictionaryPage :: Page -> Bool isDictionaryPage page = case pageTypeHeader (pageHeader page) of - DictionaryPageHeader{..} -> True + DictionaryPageHeader{} -> True _ -> False readPage :: CompressionCodec -> BS.ByteString -> IO (Maybe Page, BS.ByteString) @@ -79,25 +78,31 @@ in case fieldContents of Nothing -> (hdr, BS.drop 1 xs) - Just (remainder, elemType, identifier) -> case identifier of + Just (remainder, _elemType, identifier) -> case identifier of 1 -> let (pType, remainder') = readInt32FromBytes remainder in - readPageHeader (hdr{pageHeaderPageType = pageTypeFromInt pType}) remainder' identifier + readPageHeader + (hdr{pageHeaderPageType = pageTypeFromInt pType}) + remainder' + identifier 2 -> let - (uncompressedPageSize, remainder') = readInt32FromBytes remainder + (parsedUncompressedPageSize, remainder') = readInt32FromBytes remainder in readPageHeader - (hdr{uncompressedPageSize = uncompressedPageSize}) + (hdr{uncompressedPageSize = parsedUncompressedPageSize}) remainder' identifier 3 -> let - (compressedPageSize, remainder') = readInt32FromBytes remainder + (parsedCompressedPageSize, remainder') = readInt32FromBytes remainder in - readPageHeader (hdr{compressedPageSize = compressedPageSize}) remainder' identifier + readPageHeader + (hdr{compressedPageSize = parsedCompressedPageSize}) + remainder' + identifier 4 -> let (crc, remainder') = readInt32FromBytes remainder @@ -113,7 +118,10 @@ let (dictionaryPageHeader, remainder') = readPageTypeHeader emptyDictionaryPageHeader remainder 0 in - readPageHeader (hdr{pageTypeHeader = dictionaryPageHeader}) remainder' identifier + readPageHeader + (hdr{pageTypeHeader = dictionaryPageHeader}) + remainder' + identifier 8 -> let (dataPageHeaderV2, remainder') = readPageTypeHeader emptyDataPageHeaderV2 remainder 0 @@ -125,7 +133,7 @@ PageTypeHeader -> BS.ByteString -> Int16 -> (PageTypeHeader, BS.ByteString) readPageTypeHeader INDEX_PAGE_HEADER _ _ = error "readPageTypeHeader: unsupported INDEX_PAGE_HEADER" readPageTypeHeader PAGE_TYPE_HEADER_UNKNOWN _ _ = error "readPageTypeHeader: unsupported PAGE_TYPE_HEADER_UNKNOWN" -readPageTypeHeader hdr@(DictionaryPageHeader{..}) xs lastFieldId = +readPageTypeHeader hdr@(DictionaryPageHeader{}) xs lastFieldId = if BS.null xs then (hdr, BS.empty) else @@ -134,7 +142,7 @@ in case fieldContents of Nothing -> (hdr, BS.drop 1 xs) - Just (remainder, elemType, identifier) -> case identifier of + Just (remainder, _elemType, identifier) -> case identifier of 1 -> let (numValues, remainder') = readInt32FromBytes remainder @@ -167,7 +175,7 @@ identifier n -> error $ "readPageTypeHeader: unsupported identifier " ++ show n -readPageTypeHeader hdr@(DataPageHeader{..}) xs lastFieldId = +readPageTypeHeader hdr@(DataPageHeader{}) xs lastFieldId = if BS.null xs then (hdr, BS.empty) else @@ -176,12 +184,15 @@ in case fieldContents of Nothing -> (hdr, BS.drop 1 xs) - Just (remainder, elemType, identifier) -> case identifier of + Just (remainder, _elemType, identifier) -> case identifier of 1 -> let (numValues, remainder') = readInt32FromBytes remainder in - readPageTypeHeader (hdr{dataPageHeaderNumValues = numValues}) remainder' identifier + readPageTypeHeader + (hdr{dataPageHeaderNumValues = numValues}) + remainder' + identifier 2 -> let (enc, remainder') = readInt32FromBytes remainder @@ -212,7 +223,7 @@ in readPageTypeHeader (hdr{dataPageHeaderStatistics = stats}) remainder' identifier n -> error $ show n -readPageTypeHeader hdr@(DataPageHeaderV2{..}) xs lastFieldId = +readPageTypeHeader hdr@(DataPageHeaderV2{}) xs lastFieldId = if BS.null xs then (hdr, BS.empty) else @@ -221,22 +232,31 @@ in case fieldContents of Nothing -> (hdr, BS.drop 1 xs) - Just (remainder, elemType, identifier) -> case identifier of + Just (remainder, _elemType, identifier) -> case identifier of 1 -> let (numValues, remainder') = readInt32FromBytes remainder in - readPageTypeHeader (hdr{dataPageHeaderV2NumValues = numValues}) remainder' identifier + readPageTypeHeader + (hdr{dataPageHeaderV2NumValues = numValues}) + remainder' + identifier 2 -> let (numNulls, remainder') = readInt32FromBytes remainder in - readPageTypeHeader (hdr{dataPageHeaderV2NumNulls = numNulls}) remainder' identifier + readPageTypeHeader + (hdr{dataPageHeaderV2NumNulls = numNulls}) + remainder' + identifier 3 -> let - (numRows, remainder') = readInt32FromBytes remainder + (parsedNumRows, remainder') = readInt32FromBytes remainder in - readPageTypeHeader (hdr{dataPageHeaderV2NumRows = numRows}) remainder' identifier + readPageTypeHeader + (hdr{dataPageHeaderV2NumRows = parsedNumRows}) + remainder' + identifier 4 -> let (enc, remainder') = readInt32FromBytes remainder @@ -394,7 +414,7 @@ in case fieldContents of Nothing -> (cs, BS.drop 1 xs) - Just (remainder, elemType, identifier) -> case identifier of + Just (remainder, _elemType, identifier) -> case identifier of 1 -> let (maxInBytes, remainder') = readByteStringFromBytes remainder @@ -414,7 +434,10 @@ let (distinctCount, remainder') = readIntFromBytes @Int64 remainder in - readStatisticsFromBytes (cs{columnDistictCount = distinctCount}) remainder' identifier + readStatisticsFromBytes + (cs{columnDistictCount = distinctCount}) + remainder' + identifier 5 -> let (maxInBytes, remainder') = readByteStringFromBytes remainder diff --git a/src/DataFrame/IO/Parquet/Seeking.hs b/src/DataFrame/IO/Parquet/Seeking.hs --- a/src/DataFrame/IO/Parquet/Seeking.hs +++ b/src/DataFrame/IO/Parquet/Seeking.hs @@ -128,8 +128,8 @@ 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) 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) fRead :: (MonadIO m) => FileBufferedOrSeekable -> Stream m Word8 diff --git a/src/DataFrame/IO/Parquet/Thrift.hs b/src/DataFrame/IO/Parquet/Thrift.hs --- a/src/DataFrame/IO/Parquet/Thrift.hs +++ b/src/DataFrame/IO/Parquet/Thrift.hs @@ -68,8 +68,8 @@ let colType :: TType colType = case unsafeGetColumn colName df of - (DI.BoxedColumn _ (col :: V.Vector a)) -> haskellToTType @a - (DI.UnboxedColumn _ (col :: VU.Vector a)) -> haskellToTType @a + (DI.BoxedColumn _ (_col :: V.Vector a)) -> haskellToTType @a + (DI.UnboxedColumn _ (_col :: VU.Vector a)) -> haskellToTType @a lType = if DI.hasElemType @T.Text (unsafeGetColumn colName df) || DI.hasElemType @(Maybe T.Text) (unsafeGetColumn colName df) @@ -303,7 +303,7 @@ then return () else do let modifier = fromIntegral ((t .&. 0xf0) `shiftR` 4) :: Int16 - identifier <- + _identifier <- if modifier == 0 then readIntFromBuffer @Int16 buf pos else return 0 @@ -360,7 +360,7 @@ fieldContents <- readField metaDataBuf bufferPos lastFieldId case fieldContents of Nothing -> return metadata - Just (elemType, identifier) -> case identifier of + Just (_elemType, identifier) -> case identifier of 1 -> do parsedVersion <- readIntFromBuffer @Int32 metaDataBuf bufferPos readFileMetaData @@ -437,7 +437,8 @@ else return $ fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) let _elemType = toTType sizeAndType - parsedColumnOrders <- replicateM listSize (readColumnOrder metaDataBuf bufferPos 0) + parsedColumnOrders <- + replicateM listSize (readColumnOrder metaDataBuf bufferPos 0) readFileMetaData (metadata{columnOrders = parsedColumnOrders}) metaDataBuf @@ -469,7 +470,7 @@ fieldContents <- readField buf pos lastFieldId case fieldContents of Nothing -> return schemaElement - Just (elemType, identifier) -> case identifier of + Just (_elemType, identifier) -> case identifier of 1 -> do schemaElemType <- toIntegralType <$> readInt32FromBuffer buf pos readSchemaElement @@ -548,7 +549,7 @@ fieldContents <- readField buf pos lastFieldId case fieldContents of Nothing -> return r - Just (elemType, identifier) -> case identifier of + Just (_elemType, identifier) -> case identifier of 1 -> do sizeAndType <- readAndAdvance pos buf listSize <- @@ -587,7 +588,7 @@ fieldContents <- readField buf pos lastFieldId case fieldContents of Nothing -> return c - Just (elemType, identifier) -> case identifier of + Just (_elemType, identifier) -> case identifier of 1 -> do stringSize <- readVarIntFromBuffer @Int buf pos contents <- @@ -651,7 +652,7 @@ fieldContents <- readField buf pos lastFieldId case fieldContents of Nothing -> return cm - Just (elemType, identifier) -> case identifier of + Just (_elemType, identifier) -> case identifier of 1 -> do cType <- parquetTypeFromInt <$> readInt32FromBuffer buf pos readColumnMetadata (cm{columnType = cType}) buf pos identifier @@ -771,7 +772,7 @@ fieldContents <- readField buf pos lastFieldId case fieldContents of Nothing -> return ENCRYPTION_ALGORITHM_UNKNOWN - Just (elemType, identifier) -> case identifier of + Just (_elemType, identifier) -> case identifier of 1 -> do readAesGcmV1 ( AesGcmV1 @@ -794,7 +795,7 @@ buf pos 0 - n -> return ENCRYPTION_ALGORITHM_UNKNOWN + _n -> return ENCRYPTION_ALGORITHM_UNKNOWN readColumnOrder :: BS.ByteString -> IORef Int -> Int16 -> IO ColumnOrder @@ -802,7 +803,7 @@ fieldContents <- readField buf pos lastFieldId case fieldContents of Nothing -> return COLUMN_ORDER_UNKNOWN - Just (elemType, identifier) -> case identifier of + Just (_elemType, identifier) -> case identifier of 1 -> do -- Read begin struct and stop since this an empty struct. replicateM_ 2 (readTypeOrder buf pos 0) @@ -819,7 +820,7 @@ fieldContents <- readField buf pos lastFieldId case fieldContents of Nothing -> return v - Just (elemType, identifier) -> case identifier of + Just (_elemType, identifier) -> case identifier of 1 -> do parsedAadPrefix <- readByteString buf pos readAesGcmCtrV1 (v{aadPrefix = parsedAadPrefix}) buf pos identifier @@ -851,7 +852,7 @@ fieldContents <- readField buf pos lastFieldId case fieldContents of Nothing -> return v - Just (elemType, identifier) -> case identifier of + Just (_elemType, identifier) -> case identifier of 1 -> do parsedAadPrefix <- readByteString buf pos readAesGcmV1 (v{aadPrefix = parsedAadPrefix}) buf pos identifier @@ -886,7 +887,7 @@ fieldContents <- readField buf pos lastFieldId case fieldContents of Nothing -> return kv - Just (elemType, identifier) -> case identifier of + Just (_elemType, identifier) -> case identifier of 1 -> do k <- readString buf pos readKeyValue (kv{key = k}) buf pos identifier @@ -905,7 +906,7 @@ fieldContents <- readField buf pos lastFieldId case fieldContents of Nothing -> return pes - Just (elemType, identifier) -> case identifier of + Just (_elemType, identifier) -> case identifier of 1 -> do pType <- pageTypeFromInt <$> readInt32FromBuffer buf pos readPageEncodingStats (pes{pageEncodingPageType = pType}) buf pos identifier @@ -923,7 +924,7 @@ readParquetEncoding :: BS.ByteString -> IORef Int -> Int16 -> IO ParquetEncoding -readParquetEncoding buf pos lastFieldId = parquetEncodingFromInt <$> readInt32FromBuffer buf pos +readParquetEncoding buf pos _lastFieldId = parquetEncodingFromInt <$> readInt32FromBuffer buf pos readStatistics :: ColumnStatistics -> @@ -935,7 +936,7 @@ fieldContents <- readField buf pos lastFieldId case fieldContents of Nothing -> return cs - Just (elemType, identifier) -> case identifier of + Just (_elemType, identifier) -> case identifier of 1 -> do maxInBytes <- readByteString buf pos readStatistics (cs{columnMax = maxInBytes}) buf pos identifier @@ -984,7 +985,7 @@ fieldContents <- readField buf pos lastFieldId case fieldContents of Nothing -> return ss - Just (elemType, identifier) -> case identifier of + Just (_elemType, identifier) -> case identifier of 1 -> do parsedUnencodedByteArrayDataTypes <- readIntFromBuffer @Int64 buf pos readSizeStatistics @@ -1037,7 +1038,7 @@ fieldContents <- readField buf pos lastFieldId case fieldContents of Nothing -> pure parsedLogicalType - Just (elemType, identifier) -> case identifier of + Just (_elemType, identifier) -> case identifier of 1 -> do -- This is an empty enum and is read as a field. _ <- readField buf pos 0 @@ -1128,7 +1129,7 @@ fieldContents <- readField buf pos lastFieldId case fieldContents of Nothing -> return (DecimalType parsedPrecision parsedScale) - Just (elemType, identifier) -> case identifier of + Just (_elemType, identifier) -> case identifier of 1 -> do scale' <- readInt32FromBuffer buf pos readDecimalType parsedPrecision scale' buf pos identifier @@ -1147,7 +1148,8 @@ readTimeType parsedIsAdjustedToUTC parsedUnit buf pos lastFieldId = do fieldContents <- readField buf pos lastFieldId case fieldContents of - Nothing -> return (TimeType{isAdjustedToUTC = parsedIsAdjustedToUTC, unit = parsedUnit}) + Nothing -> + return (TimeType{isAdjustedToUTC = parsedIsAdjustedToUTC, unit = parsedUnit}) Just (elemType, identifier) -> case identifier of 1 -> do let isAdjustedToUTC' = elemType == toTType compactBooleanTrue @@ -1167,7 +1169,9 @@ readTimestampType parsedIsAdjustedToUTC parsedUnit buf pos lastFieldId = do fieldContents <- readField buf pos lastFieldId case fieldContents of - Nothing -> return (TimestampType{isAdjustedToUTC = parsedIsAdjustedToUTC, unit = parsedUnit}) + Nothing -> + return + (TimestampType{isAdjustedToUTC = parsedIsAdjustedToUTC, unit = parsedUnit}) Just (elemType, identifier) -> case identifier of 1 -> do let isAdjustedToUTC' = elemType == toTType compactBooleanTrue @@ -1182,7 +1186,7 @@ fieldContents <- readField buf pos lastFieldId case fieldContents of Nothing -> return parsedUnit - Just (elemType, identifier) -> case identifier of + Just (_elemType, identifier) -> case identifier of 1 -> do _ <- readField buf pos 0 readUnit MILLISECONDS buf pos identifier diff --git a/src/DataFrame/IO/Parquet/Time.hs b/src/DataFrame/IO/Parquet/Time.hs --- a/src/DataFrame/IO/Parquet/Time.hs +++ b/src/DataFrame/IO/Parquet/Time.hs @@ -25,7 +25,7 @@ julianDayAndNanosToUTCTime :: Integer -> Word64 -> UTCTime julianDayAndNanosToUTCTime julianDay nanosSinceMidnight = let day = julianDayToDay julianDay - secondsSinceMidnight = fromIntegral nanosSinceMidnight / 1_000_000_000 + secondsSinceMidnight = fromIntegral nanosSinceMidnight / (1_000_000_000 :: Double) diffTime = secondsToDiffTime (floor secondsSinceMidnight) in UTCTime day diffTime @@ -47,7 +47,7 @@ utcTimeToInt96 :: UTCTime -> BS.ByteString utcTimeToInt96 (UTCTime day diffTime) = let julianDay = dayToJulianDay day - nanosSinceMidnight = floor (realToFrac diffTime * 1_000_000_000) + nanosSinceMidnight = floor (realToFrac diffTime * (1_000_000_000 :: Double)) nanosBytes = word64ToLittleEndian nanosSinceMidnight julianBytes = word32ToLittleEndian (fromIntegral julianDay) in nanosBytes `BS.append` julianBytes @@ -55,7 +55,7 @@ dayToJulianDay :: Day -> Integer dayToJulianDay day = let (year, month, dayOfMonth) = toGregorian day - a = fromIntegral $ (14 - fromIntegral month) `div` 12 + a = (fromIntegral $ (14 - fromIntegral month) `div` (12 :: Integer)) :: Integer y = fromIntegral $ year + 4800 - a m = fromIntegral $ month + 12 * fromIntegral a - 3 in fromIntegral dayOfMonth diff --git a/src/DataFrame/IO/Parquet/Types.hs b/src/DataFrame/IO/Parquet/Types.hs --- a/src/DataFrame/IO/Parquet/Types.hs +++ b/src/DataFrame/IO/Parquet/Types.hs @@ -203,6 +203,7 @@ } deriving (Show, Eq) +emptyPageHeader :: PageHeader emptyPageHeader = PageHeader PAGE_TYPE_UNKNOWN 0 0 0 PAGE_TYPE_HEADER_UNKNOWN data PageTypeHeader @@ -232,7 +233,10 @@ | PAGE_TYPE_HEADER_UNKNOWN deriving (Show, Eq) +emptyDictionaryPageHeader :: PageTypeHeader emptyDictionaryPageHeader = DictionaryPageHeader 0 PARQUET_ENCODING_UNKNOWN False + +emptyDataPageHeader :: PageTypeHeader emptyDataPageHeader = DataPageHeader 0 @@ -240,6 +244,7 @@ PARQUET_ENCODING_UNKNOWN PARQUET_ENCODING_UNKNOWN emptyColumnStatistics +emptyDataPageHeaderV2 :: PageTypeHeader emptyDataPageHeaderV2 = DataPageHeaderV2 0 diff --git a/src/DataFrame/Internal/Column.hs b/src/DataFrame/Internal/Column.hs --- a/src/DataFrame/Internal/Column.hs +++ b/src/DataFrame/Internal/Column.hs @@ -28,7 +28,7 @@ import Control.DeepSeq (NFData (..), rnf) import Control.Exception (throw) import Control.Monad (forM_, when) -import Control.Monad.ST (runST) +import Control.Monad.ST (ST, runST) import Data.Bits ( complement, popCount, @@ -47,6 +47,7 @@ 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). @@ -235,10 +236,10 @@ -- | Checks if a column contains numeric values. isNumeric :: Column -> Bool -isNumeric (UnboxedColumn _ (vec :: VU.Vector a)) = case sNumeric @a of +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 +isNumeric (BoxedColumn _ (_vec :: VB.Vector a)) = case testEquality (typeRep @a) (typeRep @Integer) of Nothing -> False Just Refl -> True @@ -248,8 +249,8 @@ -} 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) + 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 @@ -450,6 +451,30 @@ [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). @@ -576,9 +601,9 @@ -- | 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 (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 +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. @@ -1088,6 +1113,33 @@ ) {-# 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. -} @@ -1238,7 +1290,7 @@ 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 [(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') @@ -1634,7 +1686,7 @@ 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) f) + STrue -> Right (VU.map (round . (realToFrac :: a -> Double)) f) SFalse -> case sIntegral @a of STrue -> Right (VU.map fromIntegral f) SFalse -> @@ -1688,3 +1740,22 @@ } ) {-# 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 --- a/src/DataFrame/Internal/DataFrame.hs +++ b/src/DataFrame/Internal/DataFrame.hs @@ -146,10 +146,10 @@ else "Nothing" get (Just (UnboxedColumn Nothing column)) = V.map (T.pack . show) (V.convert column) get Nothing = V.empty - getTextColumnFromFrame df (i, name) = get $ (V.!?) (columns d) ((M.!) (columnIndices d) name) + getTextColumnFromFrame _df (_i, name) = get $ (V.!?) (columns d) ((M.!) (columnIndices d) name) rows = transpose $ - zipWith (curry (V.toList . getTextColumnFromFrame d)) [0 ..] header + zipWith (curry (V.toList . getTextColumnFromFrame d)) [(0 :: Int) ..] header in showTable properMarkdown header types rows -- | O(1) Creates an empty dataframe diff --git a/src/DataFrame/Internal/Expression.hs b/src/DataFrame/Internal/Expression.hs --- a/src/DataFrame/Internal/Expression.hs +++ b/src/DataFrame/Internal/Expression.hs @@ -72,6 +72,7 @@ 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 @@ -234,7 +235,6 @@ (MkUnaryOp{unaryFn = atanh, unaryName = "atanh", unarySymbol = Nothing}) instance (Show a) => Show (Expr a) where - {-# HLINT ignore "Use show" #-} show :: Expr a -> String show (Col name) = "(col @" ++ show (typeRep @a) ++ " " ++ show name ++ ")" show (CastWith name tag _) = "(castWith " ++ show tag ++ " " ++ show name ++ ")" @@ -246,6 +246,7 @@ 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 @@ -267,6 +268,7 @@ 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 @@ -283,6 +285,7 @@ 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 instance (Ord a, Columnable a) => Ord (Expr a) where compare l r = compareExpr (normalize l) (normalize r) @@ -309,6 +312,7 @@ 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 :: @@ -331,6 +335,7 @@ (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 @@ -340,17 +345,19 @@ 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 (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 _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 (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 @@ -390,3 +397,4 @@ 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 new file mode 100644 --- /dev/null +++ b/src/DataFrame/Internal/Grouping.hs @@ -0,0 +1,192 @@ +{-# 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 --- a/src/DataFrame/Internal/Interpreter.hs +++ b/src/DataFrame/Internal/Interpreter.hs @@ -10,6 +10,7 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} {-# LANGUAGE UndecidableInstances #-} +{-# OPTIONS_GHC -Wno-orphans #-} module DataFrame.Internal.Interpreter ( -- * New core API @@ -31,10 +32,12 @@ 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, @@ -455,6 +458,15 @@ 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 ------------------------------------------------------------------------------- @@ -797,6 +809,29 @@ 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). @@ -875,10 +910,7 @@ addContext expr $ do v <- eval @b ctx inner case v of - Scalar _ -> - Left $ - InternalException - "Cannot apply a fold aggregation to a scalar" + Scalar x -> Right (broadcastFold ctx seed f x) Flat col -> Scalar <$> foldlColumn @b @a f seed col Group gs -> @@ -896,10 +928,14 @@ addContext expr $ do v <- eval @b ctx inner case v of - Scalar _ -> - Left $ - InternalException - "Cannot apply a merge aggregation to a scalar" + 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 -> @@ -922,17 +958,35 @@ Scalar _ -> Left $ InternalException - "Cannot apply a fold aggregation to a scalar" + "fold1 requires at least one element" Flat col -> Scalar <$> foldl1Column @a f col Group gs -> Flat . fromVector <$> V.mapM (foldl1Column @a f) gs -------------------------------------------------------------------------------- --- Aggregation helpers -------------------------------------------------------------------------------- +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. -} @@ -941,10 +995,6 @@ (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 - -------------------------------------------------------------------------------- --- Backward-compatible wrappers -------------------------------------------------------------------------------- {- | Result of interpreting an expression in a grouped context. Retained for backward compatibility with 'aggregate' and friends. diff --git a/src/DataFrame/Internal/Nullable.hs b/src/DataFrame/Internal/Nullable.hs --- a/src/DataFrame/Internal/Nullable.hs +++ b/src/DataFrame/Internal/Nullable.hs @@ -173,14 +173,14 @@ (Columnable a, Columnable (Maybe a)) => NullableArithOp (Maybe a) a (Maybe a) where - nullArithOp f Nothing _ = Nothing + 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 Nothing _ = Nothing nullCmpOp f (Just x) y = Just (f x y) -- | Non-nullable × Nullable: 'Nothing' short-circuits. @@ -191,7 +191,7 @@ ) => NullableArithOp a (Maybe a) (Maybe a) where - nullArithOp f _ Nothing = Nothing + nullArithOp _f _ Nothing = Nothing nullArithOp f x (Just y) = Just (f x y) instance @@ -202,7 +202,7 @@ ) => NullableCmpOp a (Maybe a) (Maybe Bool) where - nullCmpOp f _ Nothing = Nothing + nullCmpOp _f _ Nothing = Nothing nullCmpOp f x (Just y) = Just (f x y) -- | Nullable × Nullable: either 'Nothing' short-circuits. @@ -211,8 +211,8 @@ (Columnable a, Columnable (Maybe a)) => NullableArithOp (Maybe a) (Maybe a) (Maybe a) where - nullArithOp f Nothing _ = Nothing - nullArithOp f _ Nothing = Nothing + nullArithOp _f Nothing _ = Nothing + nullArithOp _f _ Nothing = Nothing nullArithOp f (Just x) (Just y) = Just (f x y) instance @@ -220,8 +220,8 @@ (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 Nothing _ = Nothing + nullCmpOp _f _ Nothing = Nothing nullCmpOp f (Just x) (Just y) = Just (f x y) -- --------------------------------------------------------------------------- @@ -379,8 +379,20 @@ instance NumericWidenOp Double Float where widen1 = id widen2 = realToFrac -instance NumericWidenOp Int Float where widen1 = fromIntegral; widen2 = id -instance NumericWidenOp Float Int where +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 diff --git a/src/DataFrame/Internal/Parsing.hs b/src/DataFrame/Internal/Parsing.hs --- a/src/DataFrame/Internal/Parsing.hs +++ b/src/DataFrame/Internal/Parsing.hs @@ -58,20 +58,20 @@ readInteger s = case signed decimal (T.strip s) of Left _ -> Nothing Right (value, "") -> Just value - Right (value, _) -> Nothing + 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 + 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 + Just (_value, _) -> Nothing {-# INLINE readByteStringInt #-} readByteStringDouble :: (HasCallStack) => C.ByteString -> Maybe Double @@ -82,7 +82,7 @@ case readSigned readFunc (C.strip s) of Nothing -> Nothing Just (value, "") -> Just value - Just (value, _) -> Nothing + Just (_value, _) -> Nothing {-# INLINE readByteStringDouble #-} readDouble :: (HasCallStack) => T.Text -> Maybe Double @@ -90,21 +90,21 @@ case signed double s of Left _ -> Nothing Right (value, "") -> Just value - Right (value, _) -> Nothing + 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 + 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 + Right (_value, _) -> Left s {-# INLINE readIntEither #-} readDoubleEither :: (HasCallStack) => T.Text -> Either T.Text Double @@ -112,7 +112,7 @@ case signed double s of Left _ -> Left s Right (value, "") -> Right value - Right (value, _) -> Left s + Right (_value, _) -> Left s {-# INLINE readDoubleEither #-} -- --------------------------------------------------------------------------- @@ -205,7 +205,7 @@ 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 + Fail _unconsumed ctx er -> do erpos <- hTell handle fail $ "Failed to parse CSV file around " diff --git a/src/DataFrame/Internal/Row.hs b/src/DataFrame/Internal/Row.hs --- a/src/DataFrame/Internal/Row.hs +++ b/src/DataFrame/Internal/Row.hs @@ -61,7 +61,7 @@ (!?) :: [a] -> Int -> Maybe a (!?) [] _ = Nothing (!?) (x : _) 0 = Just x -(!?) (x : xs) n = (!?) xs (n - 1) +(!?) (_x : xs) n = (!?) xs (n - 1) mkColumnFromRow :: Int -> [[Any]] -> Column mkColumnFromRow i rows = case rows of diff --git a/src/DataFrame/Internal/Statistics.hs b/src/DataFrame/Internal/Statistics.hs --- a/src/DataFrame/Internal/Statistics.hs +++ b/src/DataFrame/Internal/Statistics.hs @@ -95,15 +95,18 @@ !k = fromIntegral n' !delta = x - meanVal !meanVal' = meanVal + delta / k - !m2' = m2 + (delta ^ 2 * (k - 1)) / k - !m3' = m3 + (delta ^ 3 * (k - 1) * (k - 2)) / k ^ 2 - (3 * delta * m2) / 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) + | otherwise = (sqrt (fromIntegral n - 1) * m3) / sqrt (m2 ^ (3 :: Int)) {-# INLINE computeSkewness #-} skewness' :: (VU.Unbox a, Real a, Num a) => VU.Vector a -> Double @@ -151,8 +154,8 @@ VU.mapM ( \i -> do let !p = fromIntegral i / fromIntegral q - !position = p * fromIntegral (n - 1) - !index = floor position + !position = p * fromIntegral (n - 1) :: Double + !index = floor position :: Int !f = position - fromIntegral index x <- fmap rtf (VUM.read mutableSamp index) if f == 0 @@ -180,9 +183,9 @@ VA.sort mutableSamp V.mapM ( \i -> do - let !p = fromIntegral i / fromIntegral q + let !p = fromIntegral i / fromIntegral q :: Double !position = p * fromIntegral (n - 1) - !index = floor position + !index = floor position :: Int -- This is not exact for Ord instances. -- Figure out how to make it so. VM.read mutableSamp index @@ -201,7 +204,7 @@ meanSquaredError :: VU.Vector Double -> VU.Vector Double -> Maybe Double meanSquaredError target prediction = let - squareDiff = VU.ifoldl' (\sq i e -> (e - target VU.! i) ^ 2 + sq) 0 prediction + 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 #-} diff --git a/src/DataFrame/Lazy/IO/CSV.hs b/src/DataFrame/Lazy/IO/CSV.hs --- a/src/DataFrame/Lazy/IO/CSV.hs +++ b/src/DataFrame/Lazy/IO/CSV.hs @@ -32,11 +32,13 @@ 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) @@ -45,7 +47,12 @@ data ReadOptions = ReadOptions { hasHeader :: Bool , inferTypes :: Bool - , safeRead :: 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) @@ -53,15 +60,19 @@ , rowsRead :: !Int } -{- | By default we assume the file has a header, we infer the types on read -and we convert any rows with nullish objects into Maybe (safeRead). +{- | 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 = False + , safeRead = NoSafeRead + , safeReadOverrides = [] , rowRange = Nothing , seekPos = Nothing , totalRows = Nothing @@ -125,7 +136,11 @@ -- Freeze the mutable vectors into immutable ones nulls' <- V.unsafeFreeze nullIndices - cols <- V.mapM (freezeColumn mutableCols nulls' opts) (V.generate numColumns id) + let !columnNamesV = V.fromList columnNames + cols <- + V.mapM + (freezeColumn columnNamesV mutableCols nulls' opts) + (V.generate numColumns id) pos <- hTell handle return @@ -175,7 +190,7 @@ input' <- readIORef input unless (atEOF && input' == mempty) $ do parseWith (TIO.hGetChunk handle) (parseRow c) input' >>= \case - Fail unconsumed ctx er -> do + Fail _unconsumed ctx er -> do erpos <- hTell handle fail $ "Failed to parse CSV file around " @@ -212,15 +227,26 @@ -- | 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 mutableCols nulls opts colIndex = do +freezeColumn colNames mutableCols nulls opts colIndex = do col <- VM.unsafeRead mutableCols colIndex - frozen <- freezeColumn' (nulls V.! colIndex) col - return $! if safeRead opts then ensureOptional frozen else frozen + 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 #-} -- --------------------------------------------------------------------------- diff --git a/src/DataFrame/Lazy/Internal/Optimizer.hs b/src/DataFrame/Lazy/Internal/Optimizer.hs --- a/src/DataFrame/Lazy/Internal/Optimizer.hs +++ b/src/DataFrame/Lazy/Internal/Optimizer.hs @@ -142,7 +142,7 @@ (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) = + 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 diff --git a/src/DataFrame/Operations/Aggregation.hs b/src/DataFrame/Operations/Aggregation.hs --- a/src/DataFrame/Operations/Aggregation.hs +++ b/src/DataFrame/Operations/Aggregation.hs @@ -1,28 +1,27 @@ -{-# LANGUAGE BangPatterns #-} {-# LANGUAGE ExplicitNamespaces #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE LambdaCase #-} -{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE Strict #-} {-# LANGUAGE TypeApplications #-} -module DataFrame.Operations.Aggregation where +module DataFrame.Operations.Aggregation ( + module DataFrame.Operations.Aggregation, + 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 @@ -34,158 +33,12 @@ ) 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) - -{- | 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 = fst (dimensions 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 = fst (dimensions 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 - -{- | 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) computeRowHashes :: [Int] -> DataFrame -> VU.Vector Int computeRowHashes indices df = runST $ do diff --git a/src/DataFrame/Operations/Core.hs b/src/DataFrame/Operations/Core.hs --- a/src/DataFrame/Operations/Core.hs +++ b/src/DataFrame/Operations/Core.hs @@ -531,7 +531,7 @@ [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)) = + go acc i col@(BoxedColumn _bm (_c :: V.Vector a)) = let cname = columnName i countNulls = nulls col @@ -546,7 +546,7 @@ countNulls columnType : acc - go acc i col@(UnboxedColumn bm c) = + go acc i col@(UnboxedColumn _bm _c) = let cname = columnName i countNulls = nulls col @@ -631,7 +631,7 @@ @ -} fromUnnamedColumns :: [Column] -> DataFrame -fromUnnamedColumns = fromNamedColumns . zip (map (T.pack . show) [0 ..]) +fromUnnamedColumns = fromNamedColumns . zip (map (T.pack . show) [(0 :: Int) ..]) {- | Create a dataframe from a list of column names and rows. diff --git a/src/DataFrame/Operations/Merge.hs b/src/DataFrame/Operations/Merge.hs --- a/src/DataFrame/Operations/Merge.hs +++ b/src/DataFrame/Operations/Merge.hs @@ -1,4 +1,5 @@ {-# LANGUAGE InstanceSigs #-} +{-# OPTIONS_GHC -Wno-orphans #-} module DataFrame.Operations.Merge where diff --git a/src/DataFrame/Operations/Permutation.hs b/src/DataFrame/Operations/Permutation.hs --- a/src/DataFrame/Operations/Permutation.hs +++ b/src/DataFrame/Operations/Permutation.hs @@ -127,8 +127,8 @@ go vm n nGen VU.unsafeFreeze vm - go v (-1) _ = pure () - go v 0 _ = pure () + go _v (-1) _ = pure () + go _v 0 _ = pure () go v maxInd gen = let (n, nextGen) = randomR (1, maxInd) gen diff --git a/src/DataFrame/Operations/Statistics.hs b/src/DataFrame/Operations/Statistics.hs --- a/src/DataFrame/Operations/Statistics.hs +++ b/src/DataFrame/Operations/Statistics.hs @@ -7,6 +7,7 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} {-# LANGUAGE UndecidableInstances #-} +{-# OPTIONS_GHC -Wno-orphans #-} module DataFrame.Operations.Statistics where @@ -69,7 +70,7 @@ initDf = empty & insertVector "Statistic" (V.fromList ["Count" :: T.Text, "Percentage (%)"]) - freqs col' = + freqs _col' = L.foldl' ( \d (col'', k) -> insertVector diff --git a/src/DataFrame/Operations/Subset.hs b/src/DataFrame/Operations/Subset.hs --- a/src/DataFrame/Operations/Subset.hs +++ b/src/DataFrame/Operations/Subset.hs @@ -231,10 +231,13 @@ filterJust :: T.Text -> DataFrame -> DataFrame filterJust colName df = case getColumn colName df of Nothing -> - throw $ ColumnsNotFoundException [colName] "filterJust" (M.keys $ columnIndices df) + 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 + 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 @@ -361,7 +364,7 @@ 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))) + ++ 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) @@ -389,11 +392,11 @@ sample :: (RandomGen g) => g -> Double -> DataFrame -> DataFrame sample pureGen p df = let - rand = generateRandomVector pureGen (fst (dataframeDimensions df)) + rand = mkRandom pureGen (fst (dataframeDimensions df)) (0 :: Double) 1 cRand = col @Double "__rand__" in df - & insertUnboxedVector (name cRand) rand + & insertColumn (name cRand) rand & filterWhere (cRand .>=. Lit (1 - p)) & exclude [name cRand] @@ -410,9 +413,9 @@ (RandomGen g) => g -> Double -> DataFrame -> (DataFrame, DataFrame) randomSplit pureGen p df = let - rand = generateRandomVector pureGen (fst (dataframeDimensions df)) + rand = mkRandom pureGen (fst (dataframeDimensions df)) (0 :: Double) 1 cRand = col @Double "__rand__" - withRand = df & insertUnboxedVector (name cRand) rand + withRand = df & insertColumn (name cRand) rand in ( withRand & filterWhere (cRand .<=. Lit p) @@ -435,9 +438,9 @@ kFolds :: (RandomGen g) => g -> Int -> DataFrame -> [DataFrame] kFolds pureGen folds df = let - rand = generateRandomVector pureGen (fst (dataframeDimensions df)) + rand = mkRandom pureGen (fst (dataframeDimensions df)) (0 :: Double) 1 cRand = col @Double "__rand__" - withRand = df & insertUnboxedVector (name cRand) rand + withRand = df & insertColumn (name cRand) rand partitionSize = 1 / fromIntegral folds singleFold n d = d & filterWhere (cRand .>=. Lit (fromIntegral n * partitionSize)) @@ -450,16 +453,6 @@ d' : go (n - 1) d'' in map (exclude [name cRand]) (go (folds - 1) withRand) - -generateRandomVector :: (RandomGen g) => g -> Int -> VU.Vector Double -generateRandomVector pureGen k = VU.fromList $ go pureGen k - where - go g 0 = [] - go g n = - let - (v, g') = uniformR (0 :: Double, 1 :: Double) g - in - v : go g' (n - 1) -- | Convert any Column to a vector of Text labels (one per row). columnToTextVec :: Column -> V.Vector T.Text diff --git a/src/DataFrame/Operations/Typing.hs b/src/DataFrame/Operations/Typing.hs --- a/src/DataFrame/Operations/Typing.hs +++ b/src/DataFrame/Operations/Typing.hs @@ -1,3 +1,6 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE ScopedTypeVariables #-} @@ -8,16 +11,24 @@ 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) @@ -29,16 +40,37 @@ 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 'parseSafe' is @True@. + -- ^ 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 :: Bool - {- ^ When @True@, treat 'missingValues' and nullish strings as @Nothing@. - When @False@, only empty strings become @Nothing@. + , 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\"@). } @@ -51,12 +83,32 @@ ParseOptions { missingValues = [] , sampleSize = 100 - , parseSafe = True + , 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.map (parseDefault opts) (columns df)} +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)) = @@ -79,91 +131,199 @@ parseFromExamples :: ParseOptions -> V.Vector T.Text -> Column parseFromExamples opts cols = - let - converter = - if parseSafe opts then convertNullish (missingValues opts) else convertOnlyEmpty - examples = V.map converter (V.take (sampleSize opts) cols) - asMaybeText = V.map converter 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 - result = - case makeParsingAssumption dfmt examples of - BoolAssumption -> handleBoolAssumption asMaybeText - IntAssumption -> handleIntAssumption asMaybeText - DoubleAssumption -> handleDoubleAssumption asMaybeText - TextAssumption -> handleTextAssumption asMaybeText - DateAssumption -> handleDateAssumption dfmt asMaybeText - NoAssumption -> handleNoAssumption dfmt asMaybeText - in - if parseSafe opts then ensureOptional result else result - -handleBoolAssumption :: V.Vector (Maybe T.Text) -> Column -handleBoolAssumption asMaybeText - | parsableAsBool = - maybe (fromVector asMaybeBool) fromVector (sequenceA asMaybeBool) - | otherwise = maybe (fromVector asMaybeText) fromVector (sequenceA asMaybeText) + 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 - asMaybeBool = V.map (>>= readBool) asMaybeText - parsableAsBool = vecSameConstructor asMaybeText asMaybeBool + classify p t = if p t then Nothing else Just t -handleIntAssumption :: V.Vector (Maybe T.Text) -> Column -handleIntAssumption asMaybeText - | parsableAsInt = - maybe (fromVector asMaybeInt) fromVector (sequenceA asMaybeInt) - | parsableAsDouble = - maybe (fromVector asMaybeDouble) fromVector (sequenceA asMaybeDouble) - | otherwise = maybe (fromVector asMaybeText) fromVector (sequenceA asMaybeText) +{- | 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 - asMaybeInt = V.map (>>= readInt) asMaybeText - asMaybeDouble = V.map (>>= readDouble) asMaybeText - parsableAsInt = - vecSameConstructor asMaybeText asMaybeInt - && vecSameConstructor asMaybeText asMaybeDouble - parsableAsDouble = vecSameConstructor asMaybeText asMaybeDouble + toEither :: (T.Text -> Maybe a) -> T.Text -> Either T.Text a + toEither p t = maybe (Left t) Right (p t) -handleDoubleAssumption :: V.Vector (Maybe T.Text) -> Column -handleDoubleAssumption asMaybeText - | parsableAsDouble = - maybe (fromVector asMaybeDouble) fromVector (sequenceA asMaybeDouble) - | otherwise = maybe (fromVector asMaybeText) fromVector (sequenceA asMaybeText) - where - asMaybeDouble = V.map (>>= readDouble) asMaybeText - parsableAsDouble = vecSameConstructor asMaybeText asMaybeDouble + textToEither :: T.Text -> Either T.Text T.Text + textToEither t = if T.null t then Left t else Right t -handleDateAssumption :: DateFormat -> V.Vector (Maybe T.Text) -> Column -handleDateAssumption dateFormat asMaybeText - | parsableAsDate = - maybe (fromVector asMaybeDate) fromVector (sequenceA asMaybeDate) - | otherwise = maybe (fromVector asMaybeText) fromVector (sequenceA asMaybeText) - where - asMaybeDate = V.map (>>= parseTimeOpt dateFormat) asMaybeText - parsableAsDate = vecSameConstructor asMaybeText asMaybeDate +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 #-} -handleTextAssumption :: V.Vector (Maybe T.Text) -> Column -handleTextAssumption asMaybeText = maybe (fromVector asMaybeText) fromVector (sequenceA asMaybeText) +-- | 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 -handleNoAssumption :: DateFormat -> V.Vector (Maybe T.Text) -> Column -handleNoAssumption dateFormat asMaybeText - -- No need to check for null values. If we are in this condition, that - -- means that the examples consisted only of null values, so we can - -- confidently know that this column must be an OptionalColumn - | V.all (== Nothing) asMaybeText = fromVector asMaybeText - | parsableAsBool = fromVector asMaybeBool - | parsableAsInt = fromVector asMaybeInt - | parsableAsDouble = fromVector asMaybeDouble - | parsableAsDate = fromVector asMaybeDate - | otherwise = fromVector asMaybeText - where - asMaybeBool = V.map (>>= readBool) asMaybeText - asMaybeInt = V.map (>>= readInt) asMaybeText - asMaybeDouble = V.map (>>= readDouble) asMaybeText - asMaybeDate = V.map (>>= parseTimeOpt dateFormat) asMaybeText - parsableAsBool = vecSameConstructor asMaybeText asMaybeBool - parsableAsInt = - vecSameConstructor asMaybeText asMaybeInt - && vecSameConstructor asMaybeText asMaybeDouble - parsableAsDouble = vecSameConstructor asMaybeText asMaybeDouble - parsableAsDate = vecSameConstructor asMaybeText asMaybeDate +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 @@ -231,25 +391,46 @@ | NoAssumption | TextAssumption -parseWithTypes :: Bool -> M.Map T.Text SchemaType -> DataFrame -> DataFrame -parseWithTypes safe ts df +{- | 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 v (unsafeGetColumn k d)) d) + (\k v d -> insertColumn k (asType (resolveMode k) v (unsafeGetColumn k d)) d) df ts where - asType :: SchemaType -> Column -> Column - asType (SType (_ :: P.Proxy a)) c@(BoxedColumn _ (col :: V.Vector b)) = case typeRep @a of - App t1 t2 -> case eqTypeRep t1 (typeRep @Maybe) of + -- \| 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 + 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 @@ -258,27 +439,15 @@ Nothing -> case testEquality (typeRep @a) (typeRep @b) of Just Refl -> c Nothing -> case testEquality (typeRep @T.Text) (typeRep @b) of - Just Refl -> - if safe - then fromVector (V.map ((readMaybe @a) . T.unpack) col) - else fromVector (V.map ((read @a) . T.unpack) col) - Nothing -> - if safe - then fromVector (V.map ((readMaybe @a) . show) col) - else fromVector (V.map ((read @a) . show) col) + 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 -> - if safe - then fromVector (V.map ((readMaybe @a) . T.unpack) col) - else fromVector (V.map ((read @a) . T.unpack) col) - Nothing -> - if safe - then fromVector (V.map ((readMaybe @a) . show) col) - else fromVector (V.map ((read @a) . show) col) - asType _ c = c + 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 @@ -291,3 +460,11 @@ 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/Synthesis.hs b/src/DataFrame/Synthesis.hs --- a/src/DataFrame/Synthesis.hs +++ b/src/DataFrame/Synthesis.hs @@ -94,15 +94,15 @@ , i <- [2 .. 6] ] ++ [ p + q - | (i, p) <- zip [0 ..] existingPrograms - , (j, q) <- zip [0 ..] existingPrograms + | (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 ..] existingPrograms - , (j, q) <- zip [0 ..] existingPrograms + | (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 @@ -110,16 +110,16 @@ ++ ( if includeConds then [ F.min p q - | (i, p) <- zip [0 ..] existingPrograms - , (j, q) <- zip [0 ..] existingPrograms + | (i, p) <- zip [(0 :: Int) ..] existingPrograms + , (j, q) <- zip [(0 :: Int) ..] existingPrograms , Prelude.not (isLiteral p && isLiteral q) , Prelude.not (isConditional p || isConditional q) , p /= q , i > j ] ++ [ F.max p q - | (i, p) <- zip [0 ..] existingPrograms - , (j, q) <- zip [0 ..] existingPrograms + | (i, p) <- zip [(0 :: Int) ..] existingPrograms + , (j, q) <- zip [(0 :: Int) ..] existingPrograms , Prelude.not (isLiteral p && isLiteral q) , Prelude.not (isConditional p || isConditional q) , p /= q @@ -135,8 +135,8 @@ else [] ) ++ [ p * q - | (i, p) <- zip [0 ..] existingPrograms - , (j, q) <- zip [0 ..] existingPrograms + | (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 @@ -280,7 +280,7 @@ | otherwise = let magnitude = floor (logBase 10 (abs x)) scale = 10 ** fromIntegral (n - 1 - magnitude) - in fromIntegral (round (x * scale)) / scale + in fromIntegral (round (x * scale) :: Int) / scale roundTo2SigDigits :: Double -> Double roundTo2SigDigits = roundToSigDigits 2 @@ -326,11 +326,11 @@ MutualInformation -> ( \l r -> mutualInformationBinned - (Prelude.max 10 (ceiling (sqrt (fromIntegral (VU.length l))))) + (Prelude.max 10 (ceiling (sqrt (fromIntegral (VU.length l) :: Double)))) l r ) - PearsonCorrelation -> (\l r -> (^ 2) <$> correlation' l r) + PearsonCorrelation -> (\l r -> (^ (2 :: Int)) <$> correlation' l r) MeanSquaredError -> (\l r -> fmap negate (meanSquaredError l r)) F1 -> f1FromBinary @@ -439,7 +439,7 @@ [(Expr Bool, TypedColumn Bool)] -> [Expr Bool] pickTopNBool _ _ [] = [] -pickTopNBool df (TColumn column) ps = +pickTopNBool _df (TColumn column) ps = let l = case toVector @Double @VU.Vector column of Left e -> throw e diff --git a/src/DataFrame/Typed.hs b/src/DataFrame/Typed.hs --- a/src/DataFrame/Typed.hs +++ b/src/DataFrame/Typed.hs @@ -104,6 +104,7 @@ DataFrame.Typed.Expr.sum, mean, count, + countAll, DataFrame.Typed.Expr.minimum, DataFrame.Typed.Expr.maximum, collect, diff --git a/src/DataFrame/Typed/Expr.hs b/src/DataFrame/Typed/Expr.hs --- a/src/DataFrame/Typed/Expr.hs +++ b/src/DataFrame/Typed/Expr.hs @@ -11,6 +11,7 @@ {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE TypeOperators #-} {-# LANGUAGE UndecidableInstances #-} +{-# OPTIONS_GHC -Wno-orphans #-} {- | Type-safe expression construction for typed DataFrames. @@ -104,6 +105,7 @@ sum, mean, count, + countAll, minimum, maximum, collect, @@ -528,6 +530,10 @@ 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) diff --git a/src/DataFrame/Typed/Lazy.hs b/src/DataFrame/Typed/Lazy.hs --- a/src/DataFrame/Typed/Lazy.hs +++ b/src/DataFrame/Typed/Lazy.hs @@ -91,7 +91,7 @@ import DataFrame.Typed.Types -- | A lazy query with compile-time schema tracking. -newtype TypedLazyDataFrame (cols :: [Type]) = TLD {unTLD :: LazyDataFrame} +newtype TypedLazyDataFrame (cols :: [Type]) = TLD {_unTLD :: LazyDataFrame} instance Show (TypedLazyDataFrame cols) where show (TLD ldf) = "TypedLazyDataFrame { " ++ show ldf ++ " }" @@ -154,7 +154,7 @@ -- | A typed lazy grouped query. newtype TypedLazyGrouped (keys :: [Symbol]) (cols :: [Type]) = TLG - { unTLG :: ([T.Text], LazyDataFrame) + { _unTLG :: ([T.Text], LazyDataFrame) } -- | Group by key columns. diff --git a/tests/DecisionTree.hs b/tests/DecisionTree.hs --- a/tests/DecisionTree.hs +++ b/tests/DecisionTree.hs @@ -1,6 +1,7 @@ {-# LANGUAGE LambdaCase #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE TypeApplications #-} +{-# OPTIONS_GHC -Wno-x-partial #-} module DecisionTree where @@ -676,9 +677,9 @@ sumExpr = foldl1 (.+) (M.elems pe) case interpret @Double fixtureDF sumExpr of Left e -> assertFailure (show e) - Right (DI.TColumn col) -> - case DI.toVector @Double col of - Left e -> assertFailure (show e) + Right (DI.TColumn sumCol) -> + case DI.toVector @Double sumCol of + Left e2 -> assertFailure (show e2) Right vals -> mapM_ (\v -> assertBool ("sum ≈ 1.0, got " ++ show v) (abs (v - 1.0) < 1e-9)) @@ -695,14 +696,14 @@ Left e -> assertFailure (show e) Right (DI.TColumn hardCol) -> case DI.toVector @T.Text hardCol of - Left e -> assertFailure (show e) + Left e2 -> assertFailure (show e2) Right hardVals -> do probCols <- mapM ( \(cls, expr) -> case interpret @Double sepDF expr of - Left e -> assertFailure (show e) >> return (cls, V.empty) - Right (DI.TColumn col) -> case DI.toVector @Double col of - Left e -> assertFailure (show e) >> return (cls, V.empty) + Left e3 -> assertFailure (show e3) >> return (cls, V.empty) + Right (DI.TColumn col2) -> case DI.toVector @Double col2 of + Left e4 -> assertFailure (show e4) >> return (cls, V.empty) Right v -> return (cls, v) ) (M.toList pe) diff --git a/tests/GenDataFrame.hs b/tests/GenDataFrame.hs --- a/tests/GenDataFrame.hs +++ b/tests/GenDataFrame.hs @@ -1,5 +1,6 @@ {-# LANGUAGE GADTs #-} {-# LANGUAGE TypeApplications #-} +{-# OPTIONS_GHC -Wno-orphans #-} module GenDataFrame where @@ -39,4 +40,4 @@ instance Arbitrary DataFrame where arbitrary = genDataFrame - shrink df = [] + shrink _df = [] diff --git a/tests/IO/CSV.hs b/tests/IO/CSV.hs --- a/tests/IO/CSV.hs +++ b/tests/IO/CSV.hs @@ -10,7 +10,6 @@ import DataFrame.IO.CSV (fromCsv, fromCsvBytes) import qualified DataFrame.Internal.Column as DI import DataFrame.Internal.DataFrame (DataFrame (..), toCsv) -import qualified DataFrame.Operations.Core as D import Test.HUnit ( Test (TestCase, TestLabel), assertEqual, diff --git a/tests/Main.hs b/tests/Main.hs --- a/tests/Main.hs +++ b/tests/Main.hs @@ -1,4 +1,3 @@ -{-# LANGUAGE RecordWildCards #-} {-# LANGUAGE ScopedTypeVariables #-} module Main where @@ -34,6 +33,7 @@ import qualified Operations.Subset import qualified Operations.Take import qualified Operations.Typing +import qualified Operations.Window import qualified Operations.WriteCsv import qualified Parquet import qualified Properties @@ -62,6 +62,7 @@ ++ Operations.Subset.hunitTests ++ Operations.Take.tests ++ Operations.Typing.tests + ++ Operations.Window.tests ++ Functions.tests ++ IO.CSV.tests ++ IO.JSON.tests @@ -69,7 +70,7 @@ ++ LazyParquet.tests isSuccessful :: Result -> Bool -isSuccessful (Success{..}) = True +isSuccessful (Success{}) = True isSuccessful _ = False main :: IO () diff --git a/tests/Monad.hs b/tests/Monad.hs --- a/tests/Monad.hs +++ b/tests/Monad.hs @@ -8,7 +8,8 @@ import Test.QuickCheck import Test.QuickCheck.Monadic -roundToTwoPlaces x = fromIntegral (round (x * 100)) / 100.0 +roundToTwoPlaces :: Double -> Double +roundToTwoPlaces x = fromIntegral (round (x * 100) :: Int) / 100.0 prop_sampleM :: DataFrame -> Gen (Gen Property) prop_sampleM df = monadic' $ do @@ -23,7 +24,8 @@ let realRate = roundToTwoPlaces $ fromIntegral finalRowCount / fromIntegral rowCount let diff = abs $ expectedRate - realRate -- calculates the 99.99% confidence interval (quickcheck runs 100 tests, aim for 1/10000) - let tolerance = 3.89 * sqrt (expectedRate * (1 - expectedRate) / fromIntegral rowCount) - assert (diff <= tolerance) + let tolerance' = 3.89 * sqrt (expectedRate * (1 - expectedRate) / fromIntegral rowCount) + assert (diff <= tolerance') +tests :: [DataFrame -> Gen (Gen Property)] tests = [prop_sampleM] diff --git a/tests/Operations/Aggregations.hs b/tests/Operations/Aggregations.hs --- a/tests/Operations/Aggregations.hs +++ b/tests/Operations/Aggregations.hs @@ -44,6 +44,50 @@ ) ) +countAllAggregation :: Test +countAllAggregation = + TestCase + ( assertEqual + "countAll gives per-group row counts without a column argument" + ( D.fromNamedColumns + [ ("test1", DI.fromList [1 :: Int, 2, 3]) + , ("n", DI.fromList [6 :: Int, 3, 3]) + ] + ) + ( testData + & D.groupBy ["test1"] + & D.aggregate [F.countAll `as` "n"] + & D.sortBy [D.Asc (F.col @Int "test1")] + ) + ) + +countAllAggregationTyped :: Test +countAllAggregationTyped = + TestCase + ( assertEqual + "Typed countAll gives per-group row counts" + ( D.fromNamedColumns + [ ("test1", DI.fromList [1 :: Int, 2, 3]) + , ("n", DI.fromList [6 :: Int, 3, 3]) + ] + ) + ( testData + & either (error . show) id + . DT.freezeWithError + @[ DT.Column "test1" Int + , DT.Column "test2" Int + , DT.Column "test3" Int + , DT.Column "test4" Char + , DT.Column "test5" String + , DT.Column "test6" Integer + ] + & DT.groupBy @'["test1"] + & DT.aggregate (DT.agg @"n" DT.countAll DT.aggNil) + & DT.sortBy [DT.asc (DT.col @"test1")] + & DT.thaw + ) + ) + foldAggregationTyped :: Test foldAggregationTyped = TestCase @@ -283,6 +327,8 @@ tests :: [Test] tests = [ TestLabel "foldAggregation" foldAggregation + , TestLabel "countAllAggregation" countAllAggregation + , TestLabel "countAllAggregationTyped" countAllAggregationTyped , TestLabel "foldAggregationTyped" foldAggregationTyped , TestLabel "numericAggregation" numericAggregation , TestLabel "numericAggregationTyped" numericAggregationTyped diff --git a/tests/Operations/Apply.hs b/tests/Operations/Apply.hs --- a/tests/Operations/Apply.hs +++ b/tests/Operations/Apply.hs @@ -89,7 +89,7 @@ ( map (,D.fromList $ replicate 26 (1 :: Integer)) ["test4", "test6"] ++ -- All other fields should have their original values. - filter (\(name, col) -> name /= "test4" && name /= "test6") values + filter (\(name, _col) -> name /= "test4" && name /= "test6") values ) ) ( D.applyMany @Char diff --git a/tests/Operations/Derive.hs b/tests/Operations/Derive.hs --- a/tests/Operations/Derive.hs +++ b/tests/Operations/Derive.hs @@ -33,7 +33,7 @@ ( Just $ DI.BoxedColumn Nothing - (V.fromList (zipWith (\n c -> show n ++ [c]) [1 .. 26] ['a' .. 'z'])) + (V.fromList (zipWith (\n c -> show n ++ [c]) ([1 .. 26] :: [Int]) ['a' .. 'z'])) ) ( DI.getColumn "test4" $ D.derive @@ -52,7 +52,7 @@ TestCase ( assertEqual "typed derive works with column expression" - (zipWith (\n c -> show n ++ [c]) [1 .. 26] ['a' .. 'z']) + (zipWith (\n c -> show n ++ [c]) ([1 .. 26] :: [Int]) ['a' .. 'z']) ( DT.columnAsList @"test4" $ DT.derive @"test4" diff --git a/tests/Operations/ReadCsv.hs b/tests/Operations/ReadCsv.hs --- a/tests/Operations/ReadCsv.hs +++ b/tests/Operations/ReadCsv.hs @@ -6,8 +6,8 @@ module Operations.ReadCsv where import qualified Data.Text as T -import qualified Data.Vector.Unboxed as VU import qualified DataFrame as D +import qualified DataFrame.Lazy.IO.CSV as Lazy import DataFrame.Internal.Column (Column (..), columnTypeString) import DataFrame.Internal.DataFrame ( @@ -16,7 +16,6 @@ getColumn, ) import Test.HUnit -import Type.Reflection (typeRep) arbuthnotPath :: FilePath arbuthnotPath = "./tests/data/arbuthnot.csv" @@ -107,10 +106,441 @@ (D.fromList @T.Text ["3"]) col +-- Inference + EitherRead: first rows of "score" are clean ints, later rows are +-- "abc" / empty / "N/A". The sample is limited to the clean prefix so the +-- assumption is Int; non-parsing later rows become Left carrying raw text. +eitherReadCsvInference :: Test +eitherReadCsvInference = TestLabel "csv_eitherRead_inference" $ TestCase $ do + df <- + D.readCsvWithOpts + D.defaultReadOptions + { D.safeRead = D.EitherRead + , D.typeSpec = D.InferFromSample 5 + } + "./tests/data/unstable_csv/either_read_mixed.csv" + let expectedScore :: Column + expectedScore = + D.fromList @(Either T.Text Int) + [ Right 10 + , Right 20 + , Right 30 + , Right 40 + , Right 50 + , Left "abc" + , Left "" + , Right 80 + ] + case getColumn "score" df of + Just col -> + assertEqual + "EitherRead: 'score' = [Right 10..Right 50, Left \"abc\", Left \"\", Right 80]" + expectedScore + col + Nothing -> assertFailure "score column missing" + -- 'id' is all ints in every row, so EitherRead wraps every value as Right. + let expectedId :: Column + expectedId = + D.fromList @(Either T.Text Int) (map Right [1 .. 8]) + case getColumn "id" df of + Just col -> + assertEqual + "EitherRead: 'id' is all Right (no failures)" + expectedId + col + Nothing -> assertFailure "id column missing" + -- 'name' has a row containing "N/A"; under EitherRead the text is just + -- a value so it lands in Right. Empty cells would become Left "". + let expectedName :: Column + expectedName = + D.fromList @(Either T.Text T.Text) + (map Right ["alice", "bob", "carol", "dave", "eve", "frank", "grace", "N/A"]) + case getColumn "name" df of + Just col -> + assertEqual + "EitherRead: 'name' wraps every value as Right (no parse failures for Text)" + expectedName + col + Nothing -> assertFailure "name column missing" + +-- MaybeRead on the same CSV yields Maybe Int (bitmap-backed Int). +maybeReadCsv :: Test +maybeReadCsv = TestLabel "csv_maybeRead" $ TestCase $ do + df <- + D.readCsvWithOpts + D.defaultReadOptions{D.safeRead = D.MaybeRead} + "./tests/data/unstable_csv/either_read_mixed.csv" + -- "score" is not cleanly parseable as Int because "abc" fails; however the + -- BS inference walks all rows and falls back to Text when needed. The + -- relevant invariant for MaybeRead is that the column is nullable. + case getColumn "score" df of + Just (UnboxedColumn (Just _) _) -> pure () -- Int with bitmap + Just (BoxedColumn (Just _) _) -> pure () -- Text-backed with bitmap + Just col -> + assertFailure $ + "MaybeRead should yield a nullable column, got " + <> columnTypeString col + <> " with no bitmap" + Nothing -> assertFailure "score column missing" + +-- NoSafeRead leaves columns un-bitmapped when every row parses (id column is +-- all ints) but falls back to a Text column for mixed rows like "score". +noSafeReadCsv :: Test +noSafeReadCsv = TestLabel "csv_noSafeRead" $ TestCase $ do + df <- + D.readCsvWithOpts + D.defaultReadOptions{D.safeRead = D.NoSafeRead} + "./tests/data/unstable_csv/either_read_mixed.csv" + case getColumn "id" df of + Just (UnboxedColumn Nothing _) -> pure () -- strict Int, no bitmap + Just col -> + assertFailure $ + "NoSafeRead 'id' should be bare UnboxedColumn, got " + <> columnTypeString col + Nothing -> assertFailure "id column missing" + +-- Lazy CSV reader + EitherRead: goes through a different pipeline +-- (MutableColumn + null indices). Every row of 'id' parses cleanly, so the +-- column is @Either Text Int@ with every value in 'Right'. +lazyEitherReadCsv :: Test +lazyEitherReadCsv = TestLabel "lazy_csv_eitherRead" $ TestCase $ do + (df, _) <- + Lazy.readSeparated + ',' + Lazy.defaultOptions{Lazy.safeRead = D.EitherRead} + "./tests/data/unstable_csv/either_read_mixed.csv" + let expectedId :: Column + expectedId = + D.fromList @(Either T.Text Int) (map Right [1 .. 8]) + case getColumn "id" df of + Just col -> + assertEqual + "Lazy + EitherRead: 'id' is all Right Ints" + expectedId + col + Nothing -> assertFailure "id column missing (lazy reader)" + +-- Per-column overrides (eager): 'id' is strict Int (NoSafeRead), 'score' is +-- Either Text Int (EitherRead) recording the raw bytes for failures, 'name' +-- falls back to the default (MaybeRead) and becomes a nullable Text column. +perColumnEagerOverrides :: Test +perColumnEagerOverrides = TestLabel "csv_perColumnOverrides_eager" $ TestCase $ do + df <- + D.readCsvWithOpts + D.defaultReadOptions + { D.safeRead = D.MaybeRead + , D.safeReadOverrides = + [ ("id", D.NoSafeRead) + , ("score", D.EitherRead) + ] + , D.typeSpec = D.InferFromSample 5 + } + "./tests/data/unstable_csv/either_read_mixed.csv" + -- 'id': NoSafeRead → plain UnboxedColumn Int, no bitmap. + case getColumn "id" df of + Just (UnboxedColumn Nothing _) -> pure () + Just col -> + assertFailure $ + "override NoSafeRead for 'id' should give bare UnboxedColumn, got " + <> columnTypeString col + Nothing -> assertFailure "id column missing" + -- 'score': EitherRead → BoxedColumn of Either Text Int. + let expectedScore :: Column + expectedScore = + D.fromList @(Either T.Text Int) + [ Right 10 + , Right 20 + , Right 30 + , Right 40 + , Right 50 + , Left "abc" + , Left "" + , Right 80 + ] + case getColumn "score" df of + Just col -> + assertEqual + "override EitherRead for 'score'" + expectedScore + col + Nothing -> assertFailure "score column missing" + -- 'name': default MaybeRead kicks in → nullable column (bitmap attached). + -- Every cell is non-empty so the bitmap is all-valid, but the optional + -- wrap is still applied. + case getColumn "name" df of + Just (BoxedColumn (Just _) _) -> pure () + Just col -> + assertFailure $ + "default MaybeRead for 'name' should yield BoxedColumn with bitmap, got " + <> columnTypeString col + Nothing -> assertFailure "name column missing" + +-- Per-column overrides (lazy): same configuration as the eager test, but +-- exercises the lazy freezer's resolution of per-column modes. +perColumnLazyOverrides :: Test +perColumnLazyOverrides = TestLabel "csv_perColumnOverrides_lazy" $ TestCase $ do + (df, _) <- + Lazy.readSeparated + ',' + Lazy.defaultOptions + { Lazy.safeRead = D.MaybeRead + , Lazy.safeReadOverrides = + [ ("id", D.NoSafeRead) + , ("score", D.EitherRead) + ] + } + "./tests/data/unstable_csv/either_read_mixed.csv" + -- Lazy reader infers 'id' as Int from the first row; NoSafeRead → bare. + case getColumn "id" df of + Just (UnboxedColumn Nothing _) -> pure () + Just col -> + assertFailure $ + "lazy NoSafeRead override for 'id' → bare UnboxedColumn, got " + <> columnTypeString col + Nothing -> assertFailure "id column missing (lazy)" + -- EitherRead on 'score': every row is Right (first row was '10' → Int + -- builder), though rows that failed parsing are captured as Left via the + -- nulls list. We only assert the column is Boxed-of-Either. + case getColumn "score" df of + Just (BoxedColumn Nothing _) -> pure () + Just col -> + assertFailure $ + "lazy EitherRead override for 'score' → BoxedColumn Nothing, got " + <> columnTypeString col + Nothing -> assertFailure "score column missing (lazy)" + +-- Default=NoSafeRead, override single column 'score' → MaybeRead. +-- 'id' and 'name' keep the NoSafeRead default; 'score' becomes nullable. +overrideNoSafeReadDefaultWithMaybeRead :: Test +overrideNoSafeReadDefaultWithMaybeRead = + TestLabel "csv_override_noSafeRead_default_MaybeRead" $ TestCase $ do + df <- + D.readCsvWithOpts + D.defaultReadOptions + { D.safeRead = D.NoSafeRead + , D.safeReadOverrides = [("score", D.MaybeRead)] + } + "./tests/data/unstable_csv/either_read_mixed.csv" + -- 'id': NoSafeRead default → bare Int. + case getColumn "id" df of + Just (UnboxedColumn Nothing _) -> pure () + Just col -> + assertFailure $ + "'id' should be bare UnboxedColumn under NoSafeRead, got " + <> columnTypeString col + Nothing -> assertFailure "id column missing" + -- 'score': MaybeRead override → nullable (bitmap present). + case getColumn "score" df of + Just (UnboxedColumn (Just _) _) -> pure () -- Int with bitmap + Just (BoxedColumn (Just _) _) -> pure () -- fallback Text with bitmap + Just col -> + assertFailure $ + "'score' MaybeRead override should yield nullable column, got " + <> columnTypeString col + Nothing -> assertFailure "score column missing" + -- 'name': NoSafeRead default → Text; may have bitmap when data has + -- null cells (row 8 has "N/A" which is in missingIndicators). + -- ensureOptional is NOT forced, but a bitmap is still created when + -- the builder detects actual nulls. + case getColumn "name" df of + Just (BoxedColumn _ _) -> pure () + Just col -> + assertFailure $ + "'name' under NoSafeRead should be BoxedColumn, got " + <> columnTypeString col + Nothing -> assertFailure "name column missing" + +-- Default=NoSafeRead, override single column 'score' → EitherRead. +overrideNoSafeReadDefaultWithEitherRead :: Test +overrideNoSafeReadDefaultWithEitherRead = + TestLabel "csv_override_noSafeRead_default_EitherRead" $ TestCase $ do + df <- + D.readCsvWithOpts + D.defaultReadOptions + { D.safeRead = D.NoSafeRead + , D.safeReadOverrides = [("score", D.EitherRead)] + , D.typeSpec = D.InferFromSample 5 + } + "./tests/data/unstable_csv/either_read_mixed.csv" + -- 'score': EitherRead override → exact cell contents + let expectedScore :: Column + expectedScore = + D.fromList @(Either T.Text Int) + [ Right 10 + , Right 20 + , Right 30 + , Right 40 + , Right 50 + , Left "abc" + , Left "" + , Right 80 + ] + case getColumn "score" df of + Just col -> + assertEqual + "'score' EitherRead override from NoSafeRead default" + expectedScore + col + Nothing -> assertFailure "score column missing" + +-- Default=EitherRead, override 'id' back to NoSafeRead. +overrideEitherReadDefaultWithNoSafeRead :: Test +overrideEitherReadDefaultWithNoSafeRead = + TestLabel "csv_override_eitherRead_default_NoSafeRead" $ TestCase $ do + df <- + D.readCsvWithOpts + D.defaultReadOptions + { D.safeRead = D.EitherRead + , D.safeReadOverrides = [("id", D.NoSafeRead)] + , D.typeSpec = D.InferFromSample 5 + } + "./tests/data/unstable_csv/either_read_mixed.csv" + -- 'id': NoSafeRead override → bare Int (every row parses). + case getColumn "id" df of + Just (UnboxedColumn Nothing _) -> pure () + Just col -> + assertFailure $ + "'id' NoSafeRead override should give bare UnboxedColumn, got " + <> columnTypeString col + Nothing -> assertFailure "id column missing" + -- 'score': default EitherRead → Either Text Int + let expectedScore :: Column + expectedScore = + D.fromList @(Either T.Text Int) + [ Right 10 + , Right 20 + , Right 30 + , Right 40 + , Right 50 + , Left "abc" + , Left "" + , Right 80 + ] + case getColumn "score" df of + Just col -> + assertEqual + "'score' inherits EitherRead default" + expectedScore + col + Nothing -> assertFailure "score column missing" + -- 'name': default EitherRead → Either Text Text (all values Right). + let expectedName :: Column + expectedName = + D.fromList @(Either T.Text T.Text) + (map Right ["alice", "bob", "carol", "dave", "eve", "frank", "grace", "N/A"]) + case getColumn "name" df of + Just col -> + assertEqual + "'name' inherits EitherRead default" + expectedName + col + Nothing -> assertFailure "name column missing" + +-- Default=EitherRead, override 'name' → MaybeRead. +overrideEitherReadDefaultWithMaybeRead :: Test +overrideEitherReadDefaultWithMaybeRead = + TestLabel "csv_override_eitherRead_default_MaybeRead" $ TestCase $ do + df <- + D.readCsvWithOpts + D.defaultReadOptions + { D.safeRead = D.EitherRead + , D.safeReadOverrides = [("name", D.MaybeRead)] + , D.typeSpec = D.InferFromSample 5 + } + "./tests/data/unstable_csv/either_read_mixed.csv" + -- 'name': MaybeRead override → nullable Text column with bitmap. + case getColumn "name" df of + Just (BoxedColumn (Just _) _) -> pure () + Just col -> + assertFailure $ + "'name' MaybeRead override should yield BoxedColumn with bitmap, got " + <> columnTypeString col + Nothing -> assertFailure "name column missing" + -- 'id': default EitherRead → Either Text Int (all values Right). + let expectedId :: Column + expectedId = D.fromList @(Either T.Text Int) (map Right [1 .. 8]) + case getColumn "id" df of + Just col -> + assertEqual "'id' inherits EitherRead default" expectedId col + Nothing -> assertFailure "id column missing" + +-- Override for a column that doesn't exist in the CSV → no crash, ignored. +overrideNonExistentColumn :: Test +overrideNonExistentColumn = + TestLabel "csv_override_nonExistent" $ TestCase $ do + df <- + D.readCsvWithOpts + D.defaultReadOptions + { D.safeRead = D.NoSafeRead + , D.safeReadOverrides = [("doesNotExist", D.EitherRead)] + } + "./tests/data/unstable_csv/either_read_mixed.csv" + assertEqual "dimensions unchanged" (8, 3) (dataframeDimensions df) + +-- Empty overrides map is the same as no overrides. +emptyOverridesEqualsGlobal :: Test +emptyOverridesEqualsGlobal = + TestLabel "csv_emptyOverrides" $ TestCase $ do + df1 <- + D.readCsvWithOpts + D.defaultReadOptions{D.safeRead = D.MaybeRead} + "./tests/data/unstable_csv/either_read_mixed.csv" + df2 <- + D.readCsvWithOpts + D.defaultReadOptions + { D.safeRead = D.MaybeRead + , D.safeReadOverrides = [] + } + "./tests/data/unstable_csv/either_read_mixed.csv" + assertEqual "empty overrides produces identical DataFrame" df1 df2 + +-- Lazy path: cell-level assertions for EitherRead override on 'score'. +perColumnLazyOverridesCellLevel :: Test +perColumnLazyOverridesCellLevel = + TestLabel "csv_perColumnOverrides_lazy_cellLevel" $ TestCase $ do + (df, _) <- + Lazy.readSeparated + ',' + Lazy.defaultOptions + { Lazy.safeRead = D.NoSafeRead + , Lazy.safeReadOverrides = [("score", D.EitherRead)] + } + "./tests/data/unstable_csv/either_read_mixed.csv" + -- 'id': NoSafeRead → bare Int. + case getColumn "id" df of + Just (UnboxedColumn Nothing _) -> pure () + Just col -> + assertFailure $ + "lazy 'id' NoSafeRead → bare UnboxedColumn, got " + <> columnTypeString col + Nothing -> assertFailure "id missing (lazy)" + -- 'score': EitherRead override → BoxedColumn of Either values. + -- Lazy reader infers Int from first row; rows that fail become Left. + case getColumn "score" df of + Just (BoxedColumn Nothing _) -> pure () + Just col -> + assertFailure $ + "lazy 'score' EitherRead → BoxedColumn Nothing, got " + <> columnTypeString col + Nothing -> assertFailure "score missing (lazy)" + tests :: [Test] tests = [ specifyTypesNoInferenceFallback , specifyTypesInferFallback , specifyTypesSampleSize , testExtraFields + , eitherReadCsvInference + , maybeReadCsv + , noSafeReadCsv + , lazyEitherReadCsv + , perColumnEagerOverrides + , perColumnLazyOverrides + , -- Per-column override coverage + overrideNoSafeReadDefaultWithMaybeRead + , overrideNoSafeReadDefaultWithEitherRead + , overrideEitherReadDefaultWithNoSafeRead + , overrideEitherReadDefaultWithMaybeRead + , overrideNonExistentColumn + , emptyOverridesEqualsGlobal + , perColumnLazyOverridesCellLevel ] diff --git a/tests/Operations/Subset.hs b/tests/Operations/Subset.hs --- a/tests/Operations/Subset.hs +++ b/tests/Operations/Subset.hs @@ -187,6 +187,7 @@ , TestLabel "unit_stratifiedSplit_proportions" unit_stratifiedSplit_proportions ] +tests :: [DataFrame -> Bool] tests = [ prop_dropZero , prop_takeZero diff --git a/tests/Operations/Typing.hs b/tests/Operations/Typing.hs --- a/tests/Operations/Typing.hs +++ b/tests/Operations/Typing.hs @@ -1,5 +1,6 @@ {-# LANGUAGE DataKinds #-} {-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} module Operations.Typing where @@ -9,10 +10,11 @@ import qualified Data.Vector.Unboxed as VU import qualified DataFrame as D import qualified DataFrame.Internal.Column as DI +import DataFrame.Internal.DataFrame (getColumn) import qualified DataFrame.Operations.Typing as D import Data.Time (Day, fromGregorian) -import Test.HUnit (Test (TestCase, TestLabel), assertEqual) +import Test.HUnit (Test (TestCase, TestLabel), assertEqual, assertFailure) testData :: D.DataFrame testData = @@ -94,7 +96,7 @@ let afterParse :: [Int] afterParse = [1 .. 50] beforeParse :: [T.Text] - beforeParse = T.pack . show <$> [1 .. 50] + beforeParse = T.pack . show <$> ([1 .. 50] :: [Int]) expected = DI.fromVector $ V.fromList afterParse actual = D.parseDefault (D.defaultParseOptions{D.sampleSize = 10}) $ @@ -114,7 +116,9 @@ beforeParse :: [T.Text] beforeParse = T.pack . show - <$> [1.0 .. 50.0] ++ [3.14, 2.22, 8.55, 23.3, 12.22222235049450945049504950] + <$> ( [1.0 .. 50.0] ++ [3.14, 2.22, 8.55, 23.3, 12.22222235049450945049504950] :: + [Double] + ) expected = DI.fromVector $ V.fromList afterParse actual = D.parseDefault (D.defaultParseOptions{D.sampleSize = 10}) $ @@ -260,7 +264,7 @@ expected = DI.UnboxedColumn Nothing $ VU.fromList afterParse actual = D.parseDefault - (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = False}) + (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = D.NoSafeRead}) $ DI.fromVector $ V.fromList beforeParse in TestCase @@ -277,7 +281,7 @@ expected = DI.UnboxedColumn Nothing $ VU.fromList afterParse actual = D.parseDefault - (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = False}) + (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = D.NoSafeRead}) $ DI.fromVector $ V.fromList beforeParse in TestCase @@ -294,7 +298,7 @@ expected = DI.UnboxedColumn Nothing $ VU.fromList afterParse actual = D.parseDefault - (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = False}) + (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = D.NoSafeRead}) $ DI.fromVector $ V.fromList beforeParse in TestCase @@ -311,7 +315,7 @@ expected = DI.BoxedColumn Nothing $ V.fromList afterParse actual = D.parseDefault - (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = False}) + (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = D.NoSafeRead}) $ DI.fromVector $ V.fromList beforeParse in TestCase @@ -327,7 +331,7 @@ let expected = DI.BoxedColumn Nothing $ V.fromList texts actual = D.parseDefault - (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = False}) + (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = D.NoSafeRead}) $ DI.fromVector $ V.fromList texts assertEqual @@ -344,7 +348,7 @@ expected = DI.fromVector $ V.fromList afterParse actual = D.parseDefault - (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = False}) + (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = D.NoSafeRead}) $ DI.fromVector $ V.fromList beforeParse in TestCase @@ -362,7 +366,7 @@ expected = DI.fromVector $ V.fromList afterParse actual = D.parseDefault - (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = False}) + (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = D.NoSafeRead}) $ DI.fromVector $ V.fromList beforeParse in TestCase @@ -393,7 +397,7 @@ expected = DI.fromVector $ V.fromList afterParse actual = D.parseDefault - (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = False}) + (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = D.NoSafeRead}) $ DI.fromVector $ V.fromList beforeParse in TestCase @@ -413,7 +417,7 @@ expected = DI.fromVector $ V.fromList afterParse actual = D.parseDefault - (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = False}) + (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = D.NoSafeRead}) $ DI.fromVector $ V.fromList beforeParse in TestCase @@ -433,7 +437,7 @@ expected = DI.fromVector $ V.fromList afterParse actual = D.parseDefault - (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = False}) + (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = D.NoSafeRead}) $ DI.fromVector $ V.fromList raw assertEqual @@ -450,7 +454,7 @@ expected = DI.BoxedColumn Nothing $ V.fromList afterParse actual = D.parseDefault - (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = False}) + (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = D.NoSafeRead}) $ DI.fromVector $ V.fromList beforeParse in TestCase @@ -468,7 +472,7 @@ expected = DI.BoxedColumn Nothing $ V.fromList afterParse actual = D.parseDefault - (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = False}) + (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = D.NoSafeRead}) $ DI.fromVector $ V.fromList beforeParse in TestCase @@ -498,7 +502,7 @@ expected = DI.BoxedColumn Nothing $ V.fromList afterParse actual = D.parseDefault - (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = False}) + (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = D.NoSafeRead}) $ DI.fromVector $ V.fromList beforeParse in TestCase @@ -524,7 +528,7 @@ expected = DI.fromVector $ V.fromList afterParse actual = D.parseDefault - (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = False}) + (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = D.NoSafeRead}) $ DI.fromVector $ V.fromList beforeParse in TestCase @@ -543,7 +547,7 @@ expected = DI.fromVector $ V.fromList afterParse actual = D.parseDefault - (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = False}) + (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = D.NoSafeRead}) $ DI.fromVector $ V.fromList raw assertEqual @@ -779,6 +783,284 @@ expected actual +-- 3C. PARSING WITH SAFEREAD = EitherRead +-- +-- EitherRead wraps every column as @Either Text a@: successful parses become +-- @Right v@; failures (including nullish/empty cells) become @Left @ +-- preserving the original input verbatim. + +parseBoolsWithEitherRead :: Test +parseBoolsWithEitherRead = + let beforeParse :: [T.Text] + beforeParse = ["true", "", "false", "N/A", "True"] + afterParse :: [Either T.Text Bool] + afterParse = + [Right True, Left "", Right False, Left "N/A", Right True] + expected = DI.fromVector $ V.fromList afterParse + actual = + D.parseDefault + (D.defaultParseOptions{D.sampleSize = 5, D.parseSafe = D.EitherRead}) + $ DI.fromVector + $ V.fromList beforeParse + in TestCase + ( assertEqual + "EitherRead wraps Bools as Either Text Bool; failures carry raw text" + expected + actual + ) + +parseIntsWithEitherRead :: Test +parseIntsWithEitherRead = + -- Sample (first 5 rows) is cleanly Int; later rows exercise the Left path. + let beforeParse :: [T.Text] + beforeParse = ["1", "2", "3", "4", "42", "abc", "", "N/A"] + afterParse :: [Either T.Text Int] + afterParse = + [ Right 1 + , Right 2 + , Right 3 + , Right 4 + , Right 42 + , Left "abc" + , Left "" + , Left "N/A" + ] + expected = DI.fromVector $ V.fromList afterParse + actual = + D.parseDefault + (D.defaultParseOptions{D.sampleSize = 5, D.parseSafe = D.EitherRead}) + $ DI.fromVector + $ V.fromList beforeParse + in TestCase + ( assertEqual + "EitherRead wraps Ints as Either Text Int; failures carry raw text" + expected + actual + ) + +parseDoublesWithEitherRead :: Test +parseDoublesWithEitherRead = + -- Sample is cleanly Double; "oops" / "" come after and must surface as Left. + let beforeParse :: [T.Text] + beforeParse = ["1.5", "2.0", "3.5", "4.25", "3.14", "oops", ""] + afterParse :: [Either T.Text Double] + afterParse = + [ Right 1.5 + , Right 2.0 + , Right 3.5 + , Right 4.25 + , Right 3.14 + , Left "oops" + , Left "" + ] + expected = DI.fromVector $ V.fromList afterParse + actual = + D.parseDefault + (D.defaultParseOptions{D.sampleSize = 5, D.parseSafe = D.EitherRead}) + $ DI.fromVector + $ V.fromList beforeParse + in TestCase + ( assertEqual + "EitherRead wraps Doubles as Either Text Double; failures carry raw text" + expected + actual + ) + +parseTextsWithEitherRead :: Test +parseTextsWithEitherRead = + let beforeParse :: [T.Text] + beforeParse = ["hello", "", "world"] + afterParse :: [Either T.Text T.Text] + afterParse = [Right "hello", Left "", Right "world"] + expected = DI.fromVector $ V.fromList afterParse + actual = + D.parseDefault + (D.defaultParseOptions{D.sampleSize = 3, D.parseSafe = D.EitherRead}) + $ DI.fromVector + $ V.fromList beforeParse + in TestCase + ( assertEqual + "EitherRead wraps Texts as Either Text Text; empty cells become Left \"\"" + expected + actual + ) + +parseDatesWithEitherRead :: Test +parseDatesWithEitherRead = + -- Sample is clean dates; "not-a-date" and "" come after and become Left. + let beforeParse :: [T.Text] + beforeParse = + (T.pack . show <$> take 5 datesExpected) ++ ["not-a-date", ""] + expectedGood :: [Either T.Text Day] + expectedGood = Right <$> take 5 datesExpected + expectedBad :: [Either T.Text Day] + expectedBad = [Left "not-a-date", Left ""] + expected = DI.fromVector $ V.fromList (expectedGood ++ expectedBad) + actual = + D.parseDefault + (D.defaultParseOptions{D.sampleSize = 5, D.parseSafe = D.EitherRead}) + $ DI.fromVector + $ V.fromList beforeParse + in TestCase + ( assertEqual + "EitherRead wraps Dates as Either Text Day; failures carry raw text" + expected + actual + ) + +-- parseDefaults honours per-column overrides: 'ages' stays strict Int, +-- 'notes' is wrapped as Either Text Text. +parseDefaultsPerColumnOverrides :: Test +parseDefaultsPerColumnOverrides = TestCase $ do + let df = + D.fromNamedColumns + [ ("ages", DI.fromList @T.Text ["1", "2", "3"]) + , ("notes", DI.fromList @T.Text ["hello", "", "world"]) + ] + opts = + D.defaultParseOptions + { D.sampleSize = 3 + , D.parseSafe = D.MaybeRead + , D.parseSafeOverrides = + [ ("ages", D.NoSafeRead) + , ("notes", D.EitherRead) + ] + } + result = D.parseDefaults opts df + case getColumn "ages" result of + Just (DI.UnboxedColumn Nothing _) -> pure () -- strict Int + Just col -> + assertFailure $ + "'ages' should be bare UnboxedColumn, got " + <> show col + Nothing -> assertFailure "ages column missing" + let expectedNotes = + DI.fromList @(Either T.Text T.Text) + [Right "hello", Left "", Right "world"] + case getColumn "notes" result of + Just col -> + assertEqual + "'notes' override EitherRead → Either Text Text" + expectedNotes + col + Nothing -> assertFailure "notes column missing" + +-- Default=NoSafeRead, single override to MaybeRead. +parseDefaultsNoSafeReadWithMaybeOverride :: Test +parseDefaultsNoSafeReadWithMaybeOverride = TestCase $ do + let df = + D.fromNamedColumns + [ ("x", DI.fromList @T.Text ["1", "2", "3"]) + , ("y", DI.fromList @T.Text ["a", "", "c"]) + ] + opts = + D.defaultParseOptions + { D.sampleSize = 3 + , D.parseSafe = D.NoSafeRead + , D.parseSafeOverrides = [("y", D.MaybeRead)] + } + result = D.parseDefaults opts df + -- 'x': NoSafeRead default → bare Int, no bitmap. + case getColumn "x" result of + Just (DI.UnboxedColumn Nothing _) -> pure () + Just col -> + assertFailure $ + "'x' should be bare UnboxedColumn, got " <> show col + Nothing -> assertFailure "x column missing" + -- 'y': MaybeRead override → nullable Text with bitmap. + case getColumn "y" result of + Just (DI.BoxedColumn (Just _) _) -> pure () + Just col -> + assertFailure $ + "'y' MaybeRead override should yield nullable column, got " <> show col + Nothing -> assertFailure "y column missing" + +-- Default=EitherRead, single override to NoSafeRead. +parseDefaultsEitherReadWithNoSafeOverride :: Test +parseDefaultsEitherReadWithNoSafeOverride = TestCase $ do + let df = + D.fromNamedColumns + [ ("nums", DI.fromList @T.Text ["10", "20", "30"]) + , ("tags", DI.fromList @T.Text ["foo", "", "baz"]) + ] + opts = + D.defaultParseOptions + { D.sampleSize = 3 + , D.parseSafe = D.EitherRead + , D.parseSafeOverrides = [("nums", D.NoSafeRead)] + } + result = D.parseDefaults opts df + -- 'nums': NoSafeRead override → bare Int. + case getColumn "nums" result of + Just (DI.UnboxedColumn Nothing _) -> pure () + Just col -> + assertFailure $ + "'nums' NoSafeRead override should give bare UnboxedColumn, got " + <> show col + Nothing -> assertFailure "nums column missing" + -- 'tags': default EitherRead → Either Text Text. + let expectedTags = + DI.fromList @(Either T.Text T.Text) [Right "foo", Left "", Right "baz"] + case getColumn "tags" result of + Just col -> + assertEqual "'tags' inherits EitherRead default" expectedTags col + Nothing -> assertFailure "tags column missing" + +-- Empty overrides: behaves identically to no overrides. +parseDefaultsEmptyOverrides :: Test +parseDefaultsEmptyOverrides = TestCase $ do + let df = + D.fromNamedColumns + [("v", DI.fromList @T.Text ["1", "2", "3"])] + opts1 = D.defaultParseOptions{D.sampleSize = 3} + opts2 = opts1{D.parseSafeOverrides = []} + assertEqual + "empty parseSafeOverrides ≡ no overrides" + (D.parseDefaults opts1 df) + (D.parseDefaults opts2 df) + +-- Override for non-existent column: silently ignored, no crash. +parseDefaultsNonExistentOverride :: Test +parseDefaultsNonExistentOverride = TestCase $ do + let df = + D.fromNamedColumns + [("a", DI.fromList @T.Text ["1", "2", "3"])] + opts = + D.defaultParseOptions + { D.sampleSize = 3 + , D.parseSafeOverrides = [("z", D.EitherRead)] + } + result = D.parseDefaults opts df + -- 'a' should parse normally under the default (MaybeRead). + case getColumn "a" result of + Just (DI.UnboxedColumn (Just _) _) -> pure () + Just col -> + assertFailure $ + "'a' should be nullable UnboxedColumn under MaybeRead default, got " + <> show col + Nothing -> assertFailure "a column missing" + +parseAllNullWithEitherRead :: Test +parseAllNullWithEitherRead = + -- NoAssumption path: sample is all empty/nullish so inference can't + -- pick a numeric type. EitherRead still wraps as Either Text Text. + let beforeParse :: [T.Text] + beforeParse = ["", "", "", "N/A", ""] + afterParse :: [Either T.Text T.Text] + afterParse = [Left "", Left "", Left "", Right "N/A", Left ""] + expected = DI.fromVector $ V.fromList afterParse + actual = + D.parseDefault + (D.defaultParseOptions{D.sampleSize = 5, D.parseSafe = D.EitherRead}) + $ DI.fromVector + $ V.fromList beforeParse + in TestCase + ( assertEqual + "EitherRead with all-null sample produces Either Text Text" + expected + actual + ) + -- 4. PARSING SHOULD NOT DEPEND ON THE NUMBER OF EXAMPLES. parseBoolsWithOneExample :: Test parseBoolsWithOneExample = @@ -944,9 +1226,10 @@ ++ map Just doublesSpecial expected = DI.fromVector $ V.fromList afterParse actual = - D.parseDefault (D.defaultParseOptions{D.sampleSize = 1, D.parseSafe = False}) $ - DI.fromVector $ - V.fromList beforeParse + D.parseDefault + (D.defaultParseOptions{D.sampleSize = 1, D.parseSafe = D.NoSafeRead}) + $ DI.fromVector + $ V.fromList beforeParse in TestCase ( assertEqual "Correctly parses a mixture of Ints and Doubles as UnboxedColumn of Doubles with just one example" @@ -971,7 +1254,7 @@ expected = DI.fromVector $ V.fromList afterParse actual = D.parseDefault - (D.defaultParseOptions{D.sampleSize = 30, D.parseSafe = False}) + (D.defaultParseOptions{D.sampleSize = 30, D.parseSafe = D.NoSafeRead}) $ DI.fromVector $ V.fromList beforeParse in TestCase @@ -1001,9 +1284,10 @@ ++ map Just doublesSpecialInput expected = DI.fromVector $ V.fromList afterParse actual = - D.parseDefault (D.defaultParseOptions{D.sampleSize = 1, D.parseSafe = False}) $ - DI.fromVector $ - V.fromList beforeParse + D.parseDefault + (D.defaultParseOptions{D.sampleSize = 1, D.parseSafe = D.NoSafeRead}) + $ DI.fromVector + $ V.fromList beforeParse in TestCase ( assertEqual "Correctly parses a mixture of Ints, Doubles, empty strings, nullish as OptionalColumn of Text with just one example, when safeRead is off" @@ -1031,7 +1315,7 @@ expected = DI.fromVector $ V.fromList afterParse actual = D.parseDefault - (D.defaultParseOptions{D.sampleSize = 30, D.parseSafe = False}) + (D.defaultParseOptions{D.sampleSize = 30, D.parseSafe = D.NoSafeRead}) $ DI.fromVector $ V.fromList beforeParse in TestCase @@ -1196,6 +1480,28 @@ , TestLabel "parseTextsAndEmptyAndNullishStringsWithSafeRead" parseTextsAndEmptyAndNullishStringsWithSafeRead + , -- 3C. PARSING WITH EitherRead + TestLabel "parseBoolsWithEitherRead" parseBoolsWithEitherRead + , TestLabel "parseIntsWithEitherRead" parseIntsWithEitherRead + , TestLabel "parseDoublesWithEitherRead" parseDoublesWithEitherRead + , TestLabel "parseTextsWithEitherRead" parseTextsWithEitherRead + , TestLabel "parseDatesWithEitherRead" parseDatesWithEitherRead + , TestLabel "parseAllNullWithEitherRead" parseAllNullWithEitherRead + , TestLabel + "parseDefaultsPerColumnOverrides" + parseDefaultsPerColumnOverrides + , TestLabel + "parseDefaultsNoSafeReadWithMaybeOverride" + parseDefaultsNoSafeReadWithMaybeOverride + , TestLabel + "parseDefaultsEitherReadWithNoSafeOverride" + parseDefaultsEitherReadWithNoSafeOverride + , TestLabel + "parseDefaultsEmptyOverrides" + parseDefaultsEmptyOverrides + , TestLabel + "parseDefaultsNonExistentOverride" + parseDefaultsNonExistentOverride , -- 4. PARSING MUST NOT DEPEND ON THE NUMBER OF EXAMPLES TestLabel "parseBoolsWithOneExample" parseBoolsWithOneExample , TestLabel "parseBoolsWithManyExamples" parseBoolsWithManyExamples diff --git a/tests/Operations/Window.hs b/tests/Operations/Window.hs new file mode 100644 --- /dev/null +++ b/tests/Operations/Window.hs @@ -0,0 +1,247 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +module Operations.Window where + +import qualified Data.Text as T +import qualified DataFrame as D +import qualified DataFrame.Functions as F +import qualified DataFrame.Internal.Column as DI +import qualified DataFrame.Internal.DataFrame as DI + +import Data.Function ((&)) +import DataFrame.Operators +import Test.HUnit + +-- Similar to example discussed in https://www.sumsar.net/blog/pandas-feels-clunky-when-coming-from-r/ +purchases :: D.DataFrame +purchases = + D.fromNamedColumns + [ + ( "country" + , DI.fromList + (["US", "US", "US", "UK", "UK", "UK", "France", "France", "France"] :: [T.Text]) + ) + , ("amount", DI.fromList ([100, 200, 5000, 50, 60, 800, 30, 40, 35] :: [Double])) + , ("discount", DI.fromList ([5, 10, 0, 2, 3, 0, 1, 2, 1] :: [Double])) + ] + +globalFilterTest :: Test +globalFilterTest = + TestCase + ( assertEqual + "Global median filter removes global outliers" + 7 + ( purchases + & D.filterWhere + (F.col @Double "amount" .<=. F.median (F.col @Double "amount") * 10) + & D.nRows + ) + ) + +overMedianFilterTest :: Test +overMedianFilterTest = + TestCase + ( assertEqual + "Per-group median filter via over removes per-country outliers" + expected + actual + ) + where + actual = + purchases + & D.filterWhere + ( F.col @Double "amount" + .<=. F.over ["country"] (F.median (F.col @Double "amount")) + * 10 + ) + & D.sortBy [D.Asc (F.col @T.Text "country"), D.Asc (F.col @Double "amount")] + expected = + D.fromNamedColumns + [ + ( "country" + , DI.fromList (["France", "France", "France", "UK", "UK", "US", "US"] :: [T.Text]) + ) + , ("amount", DI.fromList ([30, 35, 40, 50, 60, 100, 200] :: [Double])) + , ("discount", DI.fromList ([1, 1, 2, 2, 3, 5, 10] :: [Double])) + ] + +globalVsOverDifferentResults :: Test +globalVsOverDifferentResults = + TestCase + ( assertBool + "Global filter and per-group over filter should produce different DataFrames" + (D.nRows globalResult /= D.nRows overResult) + ) + where + globalResult = + purchases + & D.filterWhere + (F.col @Double "amount" .<=. F.median (F.col @Double "amount") * 3) + & D.sortBy [D.Asc (F.col @T.Text "country"), D.Asc (F.col @Double "amount")] + + overResult = + purchases + & D.filterWhere + ( F.col @Double "amount" + .<=. F.over ["country"] (F.median (F.col @Double "amount")) + * 3 + ) + & D.sortBy [D.Asc (F.col @T.Text "country"), D.Asc (F.col @Double "amount")] + +overMeanDeriveTest :: Test +overMeanDeriveTest = + TestCase + ( assertEqual + "over with mean broadcasts per-group averages to all rows" + expectedMeans + actualMeans + ) + where + simpleData = + D.fromNamedColumns + [ ("group", DI.fromList (["A", "A", "B", "B", "B"] :: [T.Text])) + , ("value", DI.fromList ([10, 20, 30, 60, 90] :: [Double])) + ] + result = + simpleData + & D.derive "group_mean" (F.over ["group"] (F.mean (F.col @Double "value"))) + & D.sortBy [D.Asc (F.col @T.Text "group"), D.Asc (F.col @Double "value")] + -- A mean = 15.0, B mean = 60.0 + expectedMeans = [15.0, 15.0, 60.0, 60.0, 60.0] :: [Double] + actualMeans = case DI.getColumn "group_mean" result of + Nothing -> error "group_mean column not found" + Just col -> DI.toList @Double col + +overSumTest :: Test +overSumTest = + TestCase + ( assertEqual + "over with sum broadcasts per-group sums" + expectedSums + actualSums + ) + where + simpleData = + D.fromNamedColumns + [ ("group", DI.fromList (["X", "X", "Y", "Y"] :: [T.Text])) + , ("value", DI.fromList ([10, 20, 100, 200] :: [Int])) + ] + result = + simpleData + & D.derive "group_sum" (F.over ["group"] (F.sum (F.col @Int "value"))) + & D.sortBy [D.Asc (F.col @T.Text "group"), D.Asc (F.col @Int "value")] + expectedSums = [30, 30, 300, 300] :: [Int] + actualSums = case DI.getColumn "group_sum" result of + Nothing -> error "group_sum column not found" + Just col -> DI.toList @Int col + +overCountTest :: Test +overCountTest = + TestCase + ( assertEqual + "over with count broadcasts per-group counts" + expectedCounts + actualCounts + ) + where + simpleData = + D.fromNamedColumns + [ ("group", DI.fromList (["A", "A", "A", "B", "B"] :: [T.Text])) + , ("value", DI.fromList ([1, 2, 3, 4, 5] :: [Int])) + ] + result = + simpleData + & D.derive "group_count" (F.over ["group"] (F.count (F.col @Int "value"))) + & D.sortBy [D.Asc (F.col @T.Text "group"), D.Asc (F.col @Int "value")] + expectedCounts = [3, 3, 3, 2, 2] :: [Int] + actualCounts = case DI.getColumn "group_count" result of + Nothing -> error "group_count column not found" + Just col -> DI.toList @Int col + +mixedGlobalAndOverTest :: Test +mixedGlobalAndOverTest = + TestCase + ( assertEqual + "Can mix over (per-group) and global expressions" + expectedDeviations + actualDeviations + ) + where + simpleData = + D.fromNamedColumns + [ ("group", DI.fromList (["A", "A", "B", "B"] :: [T.Text])) + , ("value", DI.fromList ([10.0, 20.0, 100.0, 200.0] :: [Double])) + ] + result = + simpleData + & D.derive + "deviation" + (F.col @Double "value" - F.over ["group"] (F.mean (F.col @Double "value"))) + & D.sortBy [D.Asc (F.col @T.Text "group"), D.Asc (F.col @Double "value")] + expectedDeviations = [-5.0, 5.0, -50.0, 50.0] :: [Double] + actualDeviations = case DI.getColumn "deviation" result of + Nothing -> error "deviation column not found" + Just col -> DI.toList @Double col + +-- | Example from https://www.sumsar.net/blog/pandas-feels-clunky-when-coming-from-r/ +blogPostExampleGlobal :: Test +blogPostExampleGlobal = + TestCase + ( assertEqual + "Blog post example: global filter then group+aggregate" + expected + actual + ) + where + amount = F.col @Double "amount" + discount = F.col @Double "discount" + actual = + purchases + & D.filterWhere (amount .<=. F.median amount * 10) + & D.groupBy ["country"] + & D.aggregate [F.sum (amount - discount) `as` "total"] + & D.sortBy [D.Asc (F.col @T.Text "country")] + -- Global median = 60, threshold = 600 → removes 5000 and 800 + -- France: (30-1)+(40-2)+(35-1) = 101, UK: (50-2)+(60-3) = 105, US: (100-5)+(200-10) = 285 + expected = + D.fromNamedColumns + [ ("country", DI.fromList (["France", "UK", "US"] :: [T.Text])) + , ("total", DI.fromList ([101, 105, 285] :: [Double])) + ] + +blogPostExampleOver :: Test +blogPostExampleOver = + TestCase + ( assertEqual + "Blog post example: per-group filter (via over) then group+aggregate" + expected + actual + ) + where + amount = F.col @Double "amount" + discount = F.col @Double "discount" + actual = + purchases + & D.filterWhere (amount .<=. F.over ["country"] (F.median amount) * 10) + & D.groupBy ["country"] + & D.aggregate [F.sum (amount - discount) `as` "total"] + & D.sortBy [D.Asc (F.col @T.Text "country")] + expected = + D.fromNamedColumns + [ ("country", DI.fromList (["France", "UK", "US"] :: [T.Text])) + , ("total", DI.fromList ([101, 105, 285] :: [Double])) + ] + +tests :: [Test] +tests = + [ TestLabel "globalFilterTest" globalFilterTest + , TestLabel "overMedianFilterTest" overMedianFilterTest + , TestLabel "globalVsOverDifferentResults" globalVsOverDifferentResults + , TestLabel "overMeanDeriveTest" overMeanDeriveTest + , TestLabel "overSumTest" overSumTest + , TestLabel "overCountTest" overCountTest + , TestLabel "mixedGlobalAndOverTest" mixedGlobalAndOverTest + , TestLabel "blogPostExampleGlobal" blogPostExampleGlobal + , TestLabel "blogPostExampleOver" blogPostExampleOver + ] diff --git a/tests/Parquet.hs b/tests/Parquet.hs --- a/tests/Parquet.hs +++ b/tests/Parquet.hs @@ -36,10 +36,10 @@ import Test.HUnit testBothReadParquetPaths :: ((FilePath -> IO D.DataFrame) -> Test) -> Test -testBothReadParquetPaths test = +testBothReadParquetPaths mkTest = TestList - [ test D.readParquet - , test (DP._readParquetWithOpts (Just True) D.defaultParquetReadOptions) + [ mkTest D.readParquet + , mkTest (DP._readParquetWithOpts (Just True) D.defaultParquetReadOptions) ] assertColumnNullability :: diff --git a/tests/data/unstable_csv/either_read_mixed.csv b/tests/data/unstable_csv/either_read_mixed.csv new file mode 100644 --- /dev/null +++ b/tests/data/unstable_csv/either_read_mixed.csv @@ -0,0 +1,9 @@ +id,score,name +1,10,alice +2,20,bob +3,30,carol +4,40,dave +5,50,eve +6,abc,frank +7,,grace +8,80,N/A