From 31e249d060cfd2185f42227fceca61de4ffcd111 Mon Sep 17 00:00:00 2001 From: Michael Chavinda Date: Thu, 22 Oct 2026 18:32:18 -0700 Subject: [PATCH 1/3] feat!: filter takes Expr Bool; Typed mirrors DataFrame; lazy APIs take eager shapes; joinOn everywhere --- CHANGELOG.md | 40 ++- README.md | 18 +- app/Benchmark.hs | 2 +- app/LazyBenchmark.hs | 10 +- app/Synthesis.hs | 9 +- dataframe-arrow-bridge/src/DataFrame/IR.hs | 8 +- dataframe-core/src/DataFrame/Typed/Schema.hs | 10 + dataframe-csv/src/DataFrame/Typed/IO/CSV.hs | 13 + .../src/DataFrame/Fusion/Typed.hs | 5 +- dataframe-json/dataframe-json.cabal | 1 + dataframe-json/src/DataFrame/Typed/IO/JSON.hs | 30 ++ dataframe-lazy/src/DataFrame/Lazy.hs | 7 +- .../src/DataFrame/Lazy/Internal/DataFrame.hs | 72 +++-- .../src/DataFrame/Lazy/Internal/Executor.hs | 135 +++----- .../DataFrame/Lazy/Internal/LogicalPlan.hs | 13 +- .../src/DataFrame/Lazy/Internal/Optimizer.hs | 45 +-- .../DataFrame/Lazy/Internal/PhysicalPlan.hs | 9 +- dataframe-lazy/src/DataFrame/Typed/Lazy.hs | 252 ++++++++------- dataframe-lazy/test/StreamingCsv.hs | 12 +- dataframe-learn/README.md | 1 - .../DataFrame/DecisionTree/Fit.hs | 5 +- .../dataframe-operations.cabal | 1 + dataframe-operations/src/DataFrame/Monad.hs | 6 +- .../src/DataFrame/Operations/Join.hs | 15 + .../src/DataFrame/Operations/Permutation.hs | 7 + .../src/DataFrame/Operations/Subset.hs | 50 ++- .../src/DataFrame/Typed/Apply.hs | 36 ++- .../src/DataFrame/Typed/Functions.hs | 6 + .../src/DataFrame/Typed/Join.hs | 86 +++++- .../src/DataFrame/Typed/Operations.hs | 73 ++++- dataframe-parquet/src/DataFrame/IO/Parquet.hs | 2 +- .../src/DataFrame/Typed/IO/Parquet.hs | 34 +- dataframe-persistent/README.md | 2 +- .../docs/base_scripts/base_readme.md | 2 +- dataframe-persistent/tests/PersistentTests.hs | 2 +- dataframe.cabal | 9 +- docs/base/haskell_for_data_analysis.md | 15 +- docs/base_scripts/base_readme.md | 18 +- docs/coming_from_other_implementations.md | 84 ++--- docs/cookbook.md | 37 +-- docs/haskell_for_data_analysis.md | 15 +- docs/persistent_integration.md | 2 +- docs/quick_start.md | 2 +- docs/tutorial.md | 8 +- .../using_dataframe_in_a_standalone_script.md | 15 +- docs/working_with_nulls.md | 2 +- examples/Chipotle.hs | 4 +- examples/OneBillionRowChallenge.hs | 6 +- src/DataFrame.hs | 4 +- src/DataFrame/Typed.hs | 292 ++++++++++++------ src/DataFrame/Typed/Interop.hs | 63 ++++ tests/LazyParity.hs | 66 +++- tests/LazyParquet.hs | 15 +- tests/LazyProjection.hs | 12 +- tests/Main.hs | 2 + tests/Operations/Aggregations.hs | 13 +- tests/Operations/Derive.hs | 7 +- tests/Operations/Filter.hs | 13 +- tests/Operations/Join.hs | 9 +- tests/Operations/OperatorIdentity.hs | 7 +- tests/Operations/ParallelGroupBy.hs | 2 +- tests/Operations/Record.hs | 16 +- tests/Operations/Window.hs | 12 +- tests/Parity/FourWay.hs | 206 ++++++++++++ tests/Parquet.hs | 4 +- tests/Typed/Parity.hs | 106 +++++-- 66 files changed, 1454 insertions(+), 631 deletions(-) create mode 100644 dataframe-json/src/DataFrame/Typed/IO/JSON.hs create mode 100644 dataframe-operations/src/DataFrame/Typed/Functions.hs create mode 100644 src/DataFrame/Typed/Interop.hs create mode 100644 tests/Parity/FourWay.hs diff --git a/CHANGELOG.md b/CHANGELOG.md index 92e11289..04bb8abe 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,33 @@ ## Unreleased ### Breaking changes +* `filter` takes a boolean expression (`Expr Bool -> DataFrame -> DataFrame`), + matching the lazy APIs. `filterWhere` is removed; use `filter`. +* The old `filter :: Expr a -> (a -> Bool) -> …` is removed; use `filterBy`, + which takes the predicate first. The same applies to the typed API. +* `DataFrame.Monad.filterWhereM` is renamed to `filterM`. +* `DataFrame.Typed` mirrors `DataFrame`: expression combinators (`sum`, + `mean`, `count`, `lift`, `not`, casts, `asc`/`desc`, …) move to the new + `DataFrame.Typed.Functions`, imported like `DataFrame.Functions`. + `DataFrame.Typed` keeps `col`, `lit`, conditionals and operators, and now + exports the frame statistics (`mean`, `summarize`, `correlation`, …). +* `DataFrame.Lazy` takes the eager shapes: `groupBy keys` then + `aggregate aggs`; `join jt [keys]` (multi-key) plus `innerJoin`, + `leftJoin`, `rightJoin`, `fullOuterJoin`; `sortBy` takes the eager + `SortOrder` (`Asc`/`Desc` expressions). `runDataFrame` is renamed `run`. + The lazy `SortOrder (Ascending | Descending)` type is removed. + `DataFrame.Lazy` re-exports the expression operators like `DataFrame`. +* `DataFrame.Typed.Lazy` mirrors `DataFrame.Typed`: `sortBy` takes + `TSortOrder`, joins take a key list, `aggregate` takes an `as` builder, and + scans read their schema from the type (`scanCsv @Schema path`). + `fromTypedDataFrame` is removed (use `fromDataFrame`); expression + combinators come from `DataFrame.Typed.Functions`. +* Typed eager `innerJoin` and `fullOuterJoin` and lazy INNER/FULL_OUTER joins + put left columns before right ones, matching eager joins and the typed + join schemas. +* Typed `applyColumn` is renamed `apply`; typed `fromRecordsTyped` / + `toRecordsTyped` are exported from `DataFrame.Typed` as `fromRecords` / + `toRecords`. * Operators are identified by tag, not name. `UnaryOp`/`BinaryOp` expose `unaryId`/`binaryId :: OpRef …` (`Builtin` tag, `Named` text or `Anonymous` label); `unaryName`/`binaryName` are now plain functions. @@ -23,12 +50,23 @@ `DataFrame.Internal.Statistics`. ### Additions -* `DataFrame.Typed` gains `cast`, `castWithDefault`, `castEither`, +* `joinOn jt [(leftKey, rightKey)]` in all four APIs (typed: + `joinOn @'INNER @'[ '("id", "uid")]`); the result keeps the left names. +* Typed parity with the untyped API: `topBy`, `deriveMany` (built with + `as`), `join @'INNER @'["key"]`, `applyInt`, `applyDouble`, `fold`, + `empty`, `null`, `(|>)`, display/rendering (`display`, `toCsv`, + `toMarkdown`, …), and IO re-exports from `DataFrame.Typed`. +* Typed `writeParquet`, `writeParquetWithOptions`, + `readParquetFilesWithOpts`, `readSeparated`, `writeSeparated`, and the new + `DataFrame.Typed.IO.JSON` (`readJSON`, `readJSONEither`). +* `DataFrame.Typed.Functions` gains `cast`, `castWithDefault`, `castEither`, `unsafeCast` (schema-checked column), `eq`, `lt`, `gt`, `leq`, `geq`, `and`, `or`, `countDistinct`, `topK`, `bottomK`, `isNull`, `isNotNull`, `contains` and `length`, completing parity with `DataFrame.Functions`. ### Fixes +* Streaming CSV scans return columns in file order, like eager reads, + instead of alphabetical order. * A user function named like a builtin (e.g. `"add"`) no longer runs the builtin's kernel instead of the user's function. * `eqExpr` no longer equates different `lift`/`lift2` functions, `fromMaybe` diff --git a/README.md b/README.md index c6d0687c..7da8a85d 100644 --- a/README.md +++ b/README.md @@ -235,7 +235,7 @@ $(D.declareColumnsFromCsvFile "./data/housing.csv") df <- D.readCsv "./data/housing.csv" df |> D.derive "rooms_per_household" (total_rooms / households) - |> D.filterWhere (median_income .>. 5) + |> D.filter (median_income .>. 5) |> D.groupBy ["ocean_proximity"] |> D.aggregate [F.mean median_house_value `as` "avg_value"] |> D.toMarkdown' @@ -255,7 +255,7 @@ Compare this to the manual version which requires spelling out every column name -- 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) + |> D.filter (F.col @Double "median_income" .>. F.lit 5) |> D.take 5 |> D.toMarkdown' ``` @@ -336,7 +336,7 @@ D.toRecords ordersDf :: Either Text [Order] And the typed bridge — `[Order]` to `TypedDataFrame OrderSchema` and back: ```haskell -DT.thaw (DT.fromRecordsTyped xs :: DT.TypedDataFrame OrderSchema) |> D.toMarkdown' +DT.thaw (DT.fromRecords xs :: DT.TypedDataFrame OrderSchema) |> D.toMarkdown' ``` > @@ -368,7 +368,7 @@ $(D.deriveSchemaValues ''Order) orders :: IO D.DataFrame orders = do raw <- D.readCsvWithSchema orderSchema "./data/orders.csv" - pure (D.filter orderAmount (> 100) raw) + pure (D.filterBy (> 100) orderAmount raw) ``` > @@ -417,7 +417,7 @@ case DT.freeze @EmployeeSchema employees of Nothing -> "Schema mismatch!" Just tdf -> tdf |> DT.derive @"bonus" (DT.col @"salary" * DT.lit 0.1) - |> DT.filterWhere (DT.col @"salary" DT..>. DT.lit 50000) + |> DT.filter (DT.col @"salary" DT..>. DT.lit 50000) |> DT.select @'["name", "bonus"] |> DT.thaw |> D.toMarkdown' @@ -435,11 +435,11 @@ case DT.freeze @EmployeeSchema employees of ```text -- Typo in column name -> compile error -tdf |> DT.filterWhere (DT.col @"slary" DT..>. DT.lit 50000) +tdf |> DT.filter (DT.col @"slary" DT..>. DT.lit 50000) -- error: Column 'slary' not found in schema -- Wrong type -> compile error -tdf |> DT.filterWhere (DT.col @"name" DT..>. DT.lit 50000) +tdf |> DT.filter (DT.col @"name" DT..>. DT.lit 50000) -- error: Couldn't match type 'Text' with 'Double' ``` @@ -468,7 +468,7 @@ DT.thaw (DT.filterAllJust stdf |> DT.derive @"scaled" (DT.col @"score" * DT.lit ## 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: +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 operations as the eager API, and `run` passes it through an optimizer before streaming results batch-by-batch: ```haskell import qualified DataFrame.Lazy as L @@ -487,7 +487,7 @@ housingSchema = makeSchema , ("ocean_proximity", schemaType @Text) ] -lazyResult <- L.runDataFrame $ +lazyResult <- L.run $ L.scanCsv housingSchema "./data/housing.csv" |> L.filter (F.col @Double "median_income" .>. F.lit 5) |> L.derive "value_per_income" diff --git a/app/Benchmark.hs b/app/Benchmark.hs index f6867380..52d6470e 100644 --- a/app/Benchmark.hs +++ b/app/Benchmark.hs @@ -30,7 +30,7 @@ main = do let calculationTime = diffUTCTime endCalculation startCalculation putStrLn $ "Calculation Time: " ++ show calculationTime startFilter <- getCurrentTime - print $ D.filter (F.col @Double "0") (> 0.971) df D.|> D.take 10 + print $ D.filterBy (> 0.971) (F.col @Double "0") df D.|> D.take 10 endFilter <- getCurrentTime let filterTime = diffUTCTime endFilter startFilter putStrLn $ "Filter Time: " ++ show filterTime diff --git a/app/LazyBenchmark.hs b/app/LazyBenchmark.hs index 1a82c68d..55481284 100644 --- a/app/LazyBenchmark.hs +++ b/app/LazyBenchmark.hs @@ -186,18 +186,18 @@ main = do ] runQuery "Q1 — preview first 20 rows (no filter)" $ - L.runDataFrame $ + L.run $ L.take 20 $ L.scanCsv schema pathT runQuery "Q2 — filter (x > 0.999), limit 20" $ - L.runDataFrame $ + L.run $ L.take 20 $ L.filter (col @Double "x" .> lit (0.999 :: Double)) $ L.scanCsv schema pathT runQuery "Q3 — filter (x > 0.999), derive z = x*y, select [id,z], limit 20" $ - L.runDataFrame $ + L.run $ L.take 20 $ L.select ["id", "z"] $ L.derive "z" (col @Double "x" * col @Double "y") $ @@ -205,7 +205,7 @@ main = do L.scanCsv schema pathT runQuery "Q4 — filter fusion: (x > 0.5) . (y > 0.5), limit 20" $ - L.runDataFrame $ + L.run $ L.take 20 $ L.filter (col @Double "y" .> lit (0.5 :: Double)) $ L.filter (col @Double "x" .> lit (0.5 :: Double)) $ @@ -216,7 +216,7 @@ main = do ++ approx (n `div` 1000) ++ " rows expected)" ) - $ L.runDataFrame + $ L.run $ L.select ["id", "x"] $ L.filter (col @Double "x" .> lit (0.999 :: Double)) $ L.scanCsv schema pathT diff --git a/app/Synthesis.hs b/app/Synthesis.hs index 91d8d63e..f40f1076 100644 --- a/app/Synthesis.hs +++ b/app/Synthesis.hs @@ -13,6 +13,7 @@ import DataFrame.DecisionTree import DataFrame.Expression.Operators import qualified DataFrame.Functions as F import qualified DataFrame.Typed as DT +import qualified DataFrame.Typed.Functions as DTF import System.Random $( DT.deriveSchemaFromCsvFileWith @@ -106,9 +107,9 @@ clean :: ) clean tdf = tdf - |> DT.replaceColumn @"Ticket" (DT.nullLift (T.filter isAlpha) (DT.col @"Ticket")) - |> DT.replaceColumn @"Name" (DT.nullLift extractTitle (DT.col @"Name")) - |> DT.replaceColumn @"Cabin" (DT.nullLift (T.take 1) (DT.col @"Cabin")) + |> DT.replaceColumn @"Ticket" (DTF.nullLift (T.filter isAlpha) (DT.col @"Ticket")) + |> DT.replaceColumn @"Name" (DTF.nullLift extractTitle (DT.col @"Name")) + |> DT.replaceColumn @"Cabin" (DTF.nullLift (T.take 1) (DT.col @"Cabin")) |> DT.renameMany @'[ '("Name", "title") , '("Cabin", "cabin_prefix") @@ -133,7 +134,7 @@ computeAccuracy df = df |> D.select ["Survived", "prediction"] survived = DT.col @"Survived" predCol = DT.col @"prediction" - count expr = fromIntegral (DT.nRows (DT.filterWhere expr tdf)) + count expr = fromIntegral (DT.nRows (DT.filter expr tdf)) 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)) diff --git a/dataframe-arrow-bridge/src/DataFrame/IR.hs b/dataframe-arrow-bridge/src/DataFrame/IR.hs index dfc3c656..0d19ad0f 100644 --- a/dataframe-arrow-bridge/src/DataFrame/IR.hs +++ b/dataframe-arrow-bridge/src/DataFrame/IR.hs @@ -58,7 +58,7 @@ import DataFrame.Operations.Core (insertVector, renameMany) import DataFrame.Operations.Join (JoinType (..), join) import DataFrame.Operations.Permutation (SortOrder (..), sortBy) import qualified DataFrame.Operations.Statistics as Stats -import DataFrame.Operations.Subset (exclude, filterWhere, range, select) +import DataFrame.Operations.Subset (exclude, range, select) import qualified DataFrame.Operations.Subset as Subset import DataFrame.Operations.Transformations (derive) import DataFrame.Schema (Schema, makeSchema, schemaType) @@ -176,7 +176,7 @@ executePlan reader (Limit k node) = executePlan reader (Filter predJson node) = do df <- executePlan reader node case decodeExprAt @Bool predJson of - Right pred_ -> return $ filterWhere pred_ df + Right pred_ -> return $ Subset.filter pred_ df Left err -> ioError $ userError $ "DataFrame.IR.Filter: " <> err executePlan reader (Derive name exprJson node) = do df <- executePlan reader node @@ -232,10 +232,10 @@ executePlan reader (WriteCsv path node) = do return df executePlan reader (ScanCsv path schemaPairs) = do schema <- buildSchema schemaPairs - Lazy.runDataFrame (Lazy.scanCsvWith reader schema (T.pack path)) + Lazy.run (Lazy.scanCsvWith reader schema (T.pack path)) executePlan _reader (ScanParquet path schemaPairs) = do schema <- buildSchema schemaPairs - Lazy.runDataFrame (Lazy.scanParquet schema (T.pack path)) + Lazy.run (Lazy.scanParquet schema (T.pack path)) -- | Build a SortOrder from a column's runtime type. mkSortOrder :: Bool -> T.Text -> Column -> SortOrder diff --git a/dataframe-core/src/DataFrame/Typed/Schema.hs b/dataframe-core/src/DataFrame/Typed/Schema.hs index f5224b6c..0cd7ff89 100644 --- a/dataframe-core/src/DataFrame/Typed/Schema.hs +++ b/dataframe-core/src/DataFrame/Typed/Schema.hs @@ -34,6 +34,7 @@ module DataFrame.Typed.Schema ( AssertAllPresent, AssertKeyTypesMatch, AssertDisjoint, + AssertAllAbsent, AssertAllColumnsHaveType, AssertRealColumn, AllColumnsReal, @@ -201,6 +202,15 @@ type family ColumnNames (cols :: [(Symbol, Type)]) :: [Symbol] where ColumnNames '[] = '[] ColumnNames ('(n, _) ': rest) = n ': ColumnNames rest +-- | Assert that none of the new columns already exists in the schema. +type family + AssertAllAbsent (new :: [(Symbol, Type)]) (cols :: [(Symbol, Type)]) :: + Constraint + where + AssertAllAbsent '[] cols = () + AssertAllAbsent ('(n, _) ': rest) cols = + (AssertAbsent n cols, AssertAllAbsent rest cols) + -- | Assert that a column name is absent from the schema (for derive/insert). type family AssertAbsent (name :: Symbol) (cols :: [(Symbol, Type)]) :: Constraint where AssertAbsent name cols = AssertAbsentHelper name (HasName name cols) cols diff --git a/dataframe-csv/src/DataFrame/Typed/IO/CSV.hs b/dataframe-csv/src/DataFrame/Typed/IO/CSV.hs index f852f29f..375004cd 100644 --- a/dataframe-csv/src/DataFrame/Typed/IO/CSV.hs +++ b/dataframe-csv/src/DataFrame/Typed/IO/CSV.hs @@ -8,8 +8,10 @@ module DataFrame.Typed.IO.CSV ( readCsvWithError, readTsv, readCsvWithOpts, + readSeparated, writeCsv, writeTsv, + writeSeparated, ) where import Control.Applicative ((<|>)) @@ -76,3 +78,14 @@ writeCsv path = CSV.writeCsv path . thaw -- | Write a typed DataFrame to a tab-separated file. writeTsv :: FilePath -> TypedDataFrame cols -> IO () writeTsv path = CSV.writeTsv path . thaw + +-- | Read a character-separated file with custom options; throws on mismatch. +readSeparated :: + forall cols. + (KnownSchema cols, RuntimeSchema cols) => + ReadOptions -> FilePath -> IO (TypedDataFrame cols) +readSeparated = readCsvWithOpts @cols + +-- | Write a typed DataFrame with the given separator. +writeSeparated :: Char -> FilePath -> TypedDataFrame cols -> IO () +writeSeparated sep path = CSV.writeSeparated sep path . thaw diff --git a/dataframe-fusion/src/DataFrame/Fusion/Typed.hs b/dataframe-fusion/src/DataFrame/Fusion/Typed.hs index 511013ac..49ea42a1 100644 --- a/dataframe-fusion/src/DataFrame/Fusion/Typed.hs +++ b/dataframe-fusion/src/DataFrame/Fusion/Typed.hs @@ -59,7 +59,6 @@ import Prelude hiding (filter, take) import qualified DataFrame.Internal.Column as IC import qualified DataFrame.Internal.Expression as IE -import DataFrame.Lazy (SortOrder (..)) import qualified DataFrame.Typed.Aggregate as AGG import qualified DataFrame.Fusion.FFI as F @@ -218,6 +217,10 @@ joinWith how leftKey rightKey (DF leftPlan) (DF rightPlan) = do ph <- runPlanOp (F.df_plan_join lp rp cHow cOn) return (DF ph) +-- | Sort direction for 'sortBy'. +data SortOrder = Ascending | Descending + deriving (Show, Eq) + -- | Sort the result by a list of (column, direction) pairs. sortBy :: [(T.Text, SortOrder)] -> diff --git a/dataframe-json/dataframe-json.cabal b/dataframe-json/dataframe-json.cabal index 22d82dea..e58a4035 100644 --- a/dataframe-json/dataframe-json.cabal +++ b/dataframe-json/dataframe-json.cabal @@ -26,6 +26,7 @@ common warnings library import: warnings exposed-modules: DataFrame.IO.JSON + DataFrame.Typed.IO.JSON build-depends: base >= 4 && < 5, aeson >= 0.11.0.0 && < 3, bytestring >= 0.11 && < 0.14, diff --git a/dataframe-json/src/DataFrame/Typed/IO/JSON.hs b/dataframe-json/src/DataFrame/Typed/IO/JSON.hs new file mode 100644 index 00000000..bd95735b --- /dev/null +++ b/dataframe-json/src/DataFrame/Typed/IO/JSON.hs @@ -0,0 +1,30 @@ +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +-- | Typed JSON readers: decode, then validate against the schema. +module DataFrame.Typed.IO.JSON ( + readJSON, + readJSONEither, +) where + +import qualified Data.ByteString.Lazy as LBS +import qualified Data.Text as T + +import qualified DataFrame.IO.JSON as JSON +import DataFrame.Typed.Freeze (freezeOrThrow, freezeWithError) +import DataFrame.Typed.Schema (KnownSchema) +import DataFrame.Typed.Types (TypedDataFrame) + +-- | Read JSON into a typed frame (numbers decode as Maybe Double); throws. +readJSON :: + forall cols. (KnownSchema cols) => FilePath -> IO (TypedDataFrame cols) +readJSON path = JSON.readJSON path >>= freezeOrThrow @cols + +-- | Decode JSON bytes into a typed DataFrame, or explain the failure. +readJSONEither :: + forall cols. + (KnownSchema cols) => + LBS.ByteString -> Either String (TypedDataFrame cols) +readJSONEither bs = do + df <- JSON.readJSONEither bs + either (Left . T.unpack) Right (freezeWithError @cols df) diff --git a/dataframe-lazy/src/DataFrame/Lazy.hs b/dataframe-lazy/src/DataFrame/Lazy.hs index 1ce652c8..759df414 100644 --- a/dataframe-lazy/src/DataFrame/Lazy.hs +++ b/dataframe-lazy/src/DataFrame/Lazy.hs @@ -1,7 +1,12 @@ +-- | Lazy, streaming queries with the same operations as "DataFrame". module DataFrame.Lazy ( module DataFrame.Lazy.Internal.DataFrame, + module DataFrame.Expression.Operators, SortOrder (..), + JoinType (..), ) where +import DataFrame.Expression.Operators import DataFrame.Lazy.Internal.DataFrame -import DataFrame.Lazy.Internal.LogicalPlan (SortOrder (..)) +import DataFrame.Operations.Join (JoinType (..)) +import DataFrame.Operations.Permutation (SortOrder (..)) diff --git a/dataframe-lazy/src/DataFrame/Lazy/Internal/DataFrame.hs b/dataframe-lazy/src/DataFrame/Lazy/Internal/DataFrame.hs index 7e13d742..1b97d7e0 100644 --- a/dataframe-lazy/src/DataFrame/Lazy/Internal/DataFrame.hs +++ b/dataframe-lazy/src/DataFrame/Lazy/Internal/DataFrame.hs @@ -21,13 +21,13 @@ import DataFrame.Lazy.Internal.Executor (execute) import DataFrame.Lazy.Internal.LogicalPlan ( DataSource (..), LogicalPlan (..), - SortOrder (..), ) import qualified DataFrame.Lazy.Internal.Optimizer as Opt -import DataFrame.Operations.Join (JoinType) +import DataFrame.Operations.Join (JoinType (..)) +import DataFrame.Operations.Permutation (SortOrder) import DataFrame.Schema (Schema) --- | A lazy query: a 'LogicalPlan' whose execution waits for 'runDataFrame'. +-- | A lazy query: a 'LogicalPlan' whose execution waits for 'run'. data LazyDataFrame = LazyDataFrame { plan :: LogicalPlan , batchSize :: Int @@ -39,8 +39,8 @@ instance Show LazyDataFrame where <> (show (batchSize ldf) <> (", plan = " <> (show (plan ldf) <> " }"))) -- | Optimise and stream-execute the lazy query into a 'D.DataFrame'. -runDataFrame :: LazyDataFrame -> IO D.DataFrame -runDataFrame ldf = execute (Opt.optimize (batchSize ldf) (plan ldf)) +run :: LazyDataFrame -> IO D.DataFrame +run ldf = execute (Opt.optimize (batchSize ldf) (plan ldf)) -- | Lift an already-loaded eager 'D.DataFrame' into the lazy plan. fromDataFrame :: D.DataFrame -> LazyDataFrame @@ -113,37 +113,55 @@ select cols ldf = ldf{plan = Project cols (plan ldf)} filter :: E.Expr Bool -> LazyDataFrame -> LazyDataFrame filter cond ldf = ldf{plan = Filter cond (plan ldf)} --- | Join two lazy queries on the given key columns. -join :: +-- | Join two lazy queries on key columns shared by name, like eager @join@. +join :: JoinType -> [T.Text] -> LazyDataFrame -> LazyDataFrame -> LazyDataFrame +join jt keys = joinOn jt [(k, k) | k <- keys] + +-- | Join on @(leftKey, rightKey)@ pairs; the result keeps the left names. +joinOn :: JoinType -> - -- | Left join key column name - T.Text -> - -- | Right join key column name - T.Text -> - -- | Left sub-query + [(T.Text, T.Text)] -> LazyDataFrame -> - -- | Right sub-query LazyDataFrame -> LazyDataFrame -join jt leftKey rightKey left right = +joinOn jt keyPairs left right = LazyDataFrame - { plan = Join jt leftKey rightKey (plan left) (plan right) + { plan = Join jt keyPairs (plan left) (plan right) , batchSize = batchSize left } --- | Group by a set of columns and compute aggregate expressions. -groupBy :: - -- | Group-by key columns - [T.Text] -> - -- | @[(outputName, aggregateExpr)]@ - [(T.Text, E.SomeExpr)] -> - LazyDataFrame -> - LazyDataFrame -groupBy keys aggs ldf = ldf{plan = Aggregate keys aggs (plan ldf)} +-- | Inner join on the given key columns. +innerJoin :: [T.Text] -> LazyDataFrame -> LazyDataFrame -> LazyDataFrame +innerJoin = join INNER + +-- | Left join, keeping every row of the left query. +leftJoin :: [T.Text] -> LazyDataFrame -> LazyDataFrame -> LazyDataFrame +leftJoin = join LEFT + +-- | Right join, keeping every row of the right query. +rightJoin :: [T.Text] -> LazyDataFrame -> LazyDataFrame -> LazyDataFrame +rightJoin = join RIGHT + +-- | Full outer join, keeping every row of both queries. +fullOuterJoin :: [T.Text] -> LazyDataFrame -> LazyDataFrame -> LazyDataFrame +fullOuterJoin = join FULL_OUTER + +-- | A lazy query grouped by key columns, awaiting 'aggregate'. +data LazyGroupedDataFrame = LazyGroupedDataFrame [T.Text] LazyDataFrame + deriving (Show) + +-- | Group by key columns; finish with 'aggregate', as in the eager API. +groupBy :: [T.Text] -> LazyDataFrame -> LazyGroupedDataFrame +groupBy = LazyGroupedDataFrame + +-- | Compute @[(outputName, aggregateExpr)]@ for each group. +aggregate :: [E.NamedExpr] -> LazyGroupedDataFrame -> LazyDataFrame +aggregate aggs (LazyGroupedDataFrame keys ldf) = + ldf{plan = Aggregate keys aggs (plan ldf)} --- | Sort the result by the given @(column, direction)@ pairs. -sortBy :: [(T.Text, SortOrder)] -> LazyDataFrame -> LazyDataFrame -sortBy cols ldf = ldf{plan = Sort cols (plan ldf)} +-- | Sort by the given orders, as the eager @sortBy@ does. +sortBy :: [SortOrder] -> LazyDataFrame -> LazyDataFrame +sortBy ords ldf = ldf{plan = Sort ords (plan ldf)} -- | Retain at most @n@ rows. take :: Int -> LazyDataFrame -> LazyDataFrame diff --git a/dataframe-lazy/src/DataFrame/Lazy/Internal/Executor.hs b/dataframe-lazy/src/DataFrame/Lazy/Internal/Executor.hs index ac105159..e6437996 100644 --- a/dataframe-lazy/src/DataFrame/Lazy/Internal/Executor.hs +++ b/dataframe-lazy/src/DataFrame/Lazy/Internal/Executor.hs @@ -25,16 +25,13 @@ import qualified Data.ByteString as BS import qualified Data.ByteString.Char8 as C8 import qualified Data.ByteString.Unsafe as BSU import Data.IORef -import Data.Int (Int16, Int32, Int64, Int8) import qualified Data.Map as M import qualified Data.Maybe import qualified Data.Set as S import qualified Data.Text as T import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl)) -import Data.Typeable (Typeable) -import qualified Data.Vector as VB import qualified Data.Vector.Unboxed as VU -import Data.Word (Word16, Word32, Word64, Word8) +import Data.Word (Word8) import DataFrame.IO.CSV ( CsvBytesReader, CsvReader, @@ -46,7 +43,7 @@ import qualified DataFrame.Internal.Column as C import qualified DataFrame.Internal.DataFrame as D import qualified DataFrame.Internal.Expression as E import qualified DataFrame.Lazy.IO.Binary as Bin -import DataFrame.Lazy.Internal.LogicalPlan (DataSource (..), SortOrder (..)) +import DataFrame.Lazy.Internal.LogicalPlan (DataSource (..)) import DataFrame.Lazy.Internal.PhysicalPlan import qualified DataFrame.Operations.Aggregation as Agg import qualified DataFrame.Operations.Core as Core @@ -109,8 +106,8 @@ isBounded (PhysicalSort _ c) = isBounded c isBounded (PhysicalHashAggregate _ _ c) = isBounded c isBounded (PhysicalSpill c _) = isBounded c isBounded (PhysicalSourceDF _ _) = True -isBounded (PhysicalHashJoin _ _ _ l r) = isBounded l && isBounded r -isBounded (PhysicalSortMergeJoin _ _ _ l r) = isBounded l && isBounded r +isBounded (PhysicalHashJoin _ _ l r) = isBounded l && isBounded r +isBounded (PhysicalSortMergeJoin _ _ l r) = isBounded l && isBounded r -- | Execute a physical plan into a single complete 'DataFrame'. execute :: PhysicalPlan -> IO D.DataFrame @@ -149,7 +146,7 @@ buildStream (PhysicalFilter p child) = do return . Stream $ ( do mb <- pullBatch childStream - return $ fmap (Sub.filterWhere p) mb + return $ fmap (Sub.filter p) mb ) buildStream (PhysicalProject cols child) = do childStream <- buildStream child @@ -182,10 +179,9 @@ buildStream (PhysicalLimit n child) = do modifyIORef' countRef (+ toTake) return $ Just (Sub.take toTake df) ) -buildStream (PhysicalSort cols child) = do +buildStream (PhysicalSort ords child) = do df <- execute child - let sortOrds = fmap (toPermSortOrder df) cols - materialized (Perm.sortBy sortOrds df) + materialized (Perm.sortBy ords df) buildStream (PhysicalHashAggregate keys aggs child) | isBounded child = do df <- execute child @@ -230,29 +226,26 @@ buildStream (PhysicalSourceDF bs df) = do batch = Sub.range (i, i + n) df writeIORef posRef (i + n) return (Just batch) -buildStream (PhysicalHashJoin jt leftKey rightKey leftPlan rightPlan) +buildStream (PhysicalHashJoin jt keyPairs leftPlan rightPlan) | isBounded leftPlan = do leftDf <- execute leftPlan rightDf <- execute rightPlan - materialized (performJoin jt leftKey rightKey leftDf rightDf) -buildStream (PhysicalHashJoin jt leftKey rightKey leftPlan rightPlan) = + materialized (performJoin jt keyPairs leftDf rightDf) +buildStream (PhysicalHashJoin jt keyPairs leftPlan rightPlan) = case jt of Join.INNER -> streamingHashJoin assembleInnerBatch Join.LEFT -> streamingHashJoin assembleLeftBatch _ -> do leftDf <- execute leftPlan rightDf <- execute rightPlan - materialized (performJoin jt leftKey rightKey leftDf rightDf) + materialized (performJoin jt keyPairs leftDf rightDf) where streamingHashJoin assembleFn = do rightDf <- execute rightPlan - let rightDf' = - if leftKey == rightKey - then rightDf - else Core.rename rightKey leftKey rightDf - joinKey = leftKey - csSet = S.fromList [joinKey] - rightHashes = Join.buildHashColumn [joinKey] rightDf' + let rightDf' = renameRightKeys keyPairs rightDf + joinKeys = map fst keyPairs + csSet = S.fromList joinKeys + rightHashes = Join.buildHashColumn joinKeys rightDf' ci = Join.buildCompactIndex rightHashes leftStream <- buildStream leftPlan return . Stream $ do @@ -260,7 +253,7 @@ buildStream (PhysicalHashJoin jt leftKey rightKey leftPlan rightPlan) = case mBatch of Nothing -> return Nothing Just probeBatch -> do - let probeHashes = Join.buildHashColumn [joinKey] probeBatch + let probeHashes = Join.buildHashColumn joinKeys probeBatch (probeIxs, buildIxs) = Join.hashProbeKernel ci probeHashes return . Just $ assembleFn csSet probeBatch rightDf' probeIxs buildIxs @@ -277,10 +270,10 @@ buildStream (PhysicalHashJoin jt leftKey rightKey leftPlan rightPlan) = in Join.assembleLeft csSet probeBatch rightDf' allProbeIxs allBuildIxs assembleInnerBatch = Join.assembleInner -buildStream (PhysicalSortMergeJoin jt leftKey rightKey leftPlan rightPlan) = do +buildStream (PhysicalSortMergeJoin jt keyPairs leftPlan rightPlan) = do leftDf <- execute leftPlan rightDf <- execute rightPlan - materialized (performJoin jt leftKey rightKey leftDf rightDf) + materialized (performJoin jt keyPairs leftDf rightDf) -- | One worker's loop: pull shared batches into a per-worker accumulator. workerLoop :: @@ -500,7 +493,7 @@ executeCsvScan path sep reader cfg = do Just df -> let df' = case scanPushdownPredicate cfg of Nothing -> df - Just p -> Sub.filterWhere p df + Just p -> Sub.filter p df in return (Just df') ) @@ -521,7 +514,7 @@ executeCsvScanStreamingWith :: ScanConfig -> IO Stream executeCsvScanStreamingWith path sep reader cfg = do - let opts = scanReadOptions sep cfg + let schemaOpts = scanReadOptions sep cfg batchSz = scanBatchSize cfg windowBytes = 64 * 1024 * 1024 :: Int queue <- newTBQueueIO 8 @@ -530,6 +523,11 @@ executeCsvScanStreamingWith path sep reader cfg = do ( withFile path ReadMode $ \h -> do rawHeader <- C8.hGetLine h let header = stripUtf8Bom rawHeader + opts = + schemaOpts + { readColumns = + fileOrder (headerNames header) <$> readColumns schemaOpts + } let feed bytes = unless (BS.null bytes) $ do let window = header <> BS.singleton nl <> bytes @@ -568,7 +566,7 @@ executeCsvScanStreamingWith path sep reader cfg = do Right (Just df) -> let df' = case scanPushdownPredicate cfg of Nothing -> df - Just p -> Sub.filterWhere p df + Just p -> Sub.filter p df in return (Just df') where nl :: Word8 @@ -577,6 +575,17 @@ executeCsvScanStreamingWith path sep reader cfg = do stripUtf8Bom bytes = Data.Maybe.fromMaybe bytes (BS.stripPrefix "\xEF\xBB\xBF" bytes) + headerNames = + map (T.dropAround (\c -> c == '"' || c == '\r') . T.pack . C8.unpack) + . C8.split sep + + -- File order, like eager reads; unknown names trail so they still error. + fileOrder names requested = + let wanted = S.fromList requested + present = S.fromList names + in filter (`S.member` wanted) names + <> filter (`S.notMember` present) requested + -- | Find the last record-ending LF in a CSV buffer. lastCompleteRecordNewline :: BS.ByteString -> Maybe Int lastCompleteRecordNewline bytes @@ -641,61 +650,17 @@ splitCsvAtNewlines n path = do Just i -> off + i + 1 Nothing -> BS.length body --- | Route a join to 'Operations.Join', renaming the right key if needed. +-- | Join as the eager API does, after renaming right keys to left names. performJoin :: - Join.JoinType -> T.Text -> T.Text -> D.DataFrame -> D.DataFrame -> D.DataFrame -performJoin jt leftKey rightKey leftDf rightDf = - case jt of - Join.LEFT -> Join.join jt [leftKey] leftDf rightRenamed - Join.RIGHT -> Join.join jt [leftKey] leftDf rightRenamed - _ -> Join.join jt [leftKey] rightRenamed leftDf - where - rightRenamed - | leftKey == rightKey = rightDf - | otherwise = Core.rename rightKey leftKey rightDf - --- | Convert a plan-level (column, direction) into a Permutation 'SortOrder'. -toPermSortOrder :: D.DataFrame -> (T.Text, SortOrder) -> Perm.SortOrder -toPermSortOrder df (col, dir) = - case M.lookup col (D.columnIndices df) of - Nothing -> mk @T.Text - Just idx -> dispatch (D.columns df VB.! idx) - where - mk :: forall a. (C.Columnable a, Ord a) => Perm.SortOrder - mk = case dir of - Ascending -> Perm.Asc (E.Col @a col) - Descending -> Perm.Desc (E.Col @a col) - - dispatch :: C.Column -> Perm.SortOrder - dispatch column = case column of - c@(C.MergedColumn _ _) -> dispatch (C.materializeMerged c) - C.PackedText{} -> mk @T.Text - C.BoxedColumn _ (_ :: VB.Vector b) -> pick @b - C.UnboxedColumn _ (_ :: VU.Vector b) -> pick @b - - pick :: forall b. (Typeable b) => Perm.SortOrder - pick = - tryT @b @Int $ - tryT @b @Double $ - tryT @b @Float $ - tryT @b @Integer $ - tryT @b @Int8 $ - tryT @b @Int16 $ - tryT @b @Int32 $ - tryT @b @Int64 $ - tryT @b @Word8 $ - tryT @b @Word16 $ - tryT @b @Word32 $ - tryT @b @Word64 $ - tryT @b @Bool $ - tryT @b @Char $ - tryT @b @T.Text (mk @T.Text) - - tryT :: - forall b a. - (Typeable b, C.Columnable a, Ord a) => - Perm.SortOrder -> - Perm.SortOrder - tryT fallback = case testEquality (typeRep @b) (typeRep @a) of - Just Refl -> mk @a - Nothing -> fallback + Join.JoinType -> + [(T.Text, T.Text)] -> + D.DataFrame -> + D.DataFrame -> + D.DataFrame +performJoin jt keyPairs leftDf rightDf = + Join.join jt (map fst keyPairs) leftDf (renameRightKeys keyPairs rightDf) + +-- | Give each right key its left counterpart's name. +renameRightKeys :: [(T.Text, T.Text)] -> D.DataFrame -> D.DataFrame +renameRightKeys keyPairs = + Core.renameMany [(r, l) | (l, r) <- keyPairs, l /= r] diff --git a/dataframe-lazy/src/DataFrame/Lazy/Internal/LogicalPlan.hs b/dataframe-lazy/src/DataFrame/Lazy/Internal/LogicalPlan.hs index 4f7c0968..906fb022 100644 --- a/dataframe-lazy/src/DataFrame/Lazy/Internal/LogicalPlan.hs +++ b/dataframe-lazy/src/DataFrame/Lazy/Internal/LogicalPlan.hs @@ -7,6 +7,7 @@ import DataFrame.IO.CSV (CsvBytesReader, CsvReader) import qualified DataFrame.Internal.DataFrame as D import qualified DataFrame.Internal.Expression as E import DataFrame.Operations.Join (JoinType) +import DataFrame.Operations.Permutation (SortOrder) import DataFrame.Schema (Schema) -- | Data source for a scan node. @@ -26,10 +27,6 @@ instance Show DataSource where "CsvSourceStreamingBytes " ++ show path ++ " " ++ show sep ++ " " show (ParquetSource path) = "ParquetSource " ++ show path --- | Sort direction used in Sort nodes and the public API. -data SortOrder = Ascending | Descending - deriving (Show, Eq, Ord) - -- | Relational-algebra tree describing what the query computes. data LogicalPlan = -- | Read columns described by the schema from a source. @@ -40,12 +37,12 @@ data LogicalPlan Filter (E.Expr Bool) LogicalPlan | -- | Add or overwrite a column via an expression. Derive T.Text E.SomeExpr LogicalPlan - | -- | Join two sub-plans on the given key columns. - Join JoinType T.Text T.Text LogicalPlan LogicalPlan + | -- | Join two sub-plans on @(leftKey, rightKey)@ pairs. + Join JoinType [(T.Text, T.Text)] LogicalPlan LogicalPlan | -- | Group then aggregate. Aggregate [T.Text] [(T.Text, E.SomeExpr)] LogicalPlan - | -- | Sort by a list of (column, direction) pairs. - Sort [(T.Text, SortOrder)] LogicalPlan + | -- | Sort by the given orders, as the eager @sortBy@ does. + Sort [SortOrder] LogicalPlan | -- | Retain at most N rows. Limit Int LogicalPlan | -- | Lift an already-loaded DataFrame into the lazy plan. diff --git a/dataframe-lazy/src/DataFrame/Lazy/Internal/Optimizer.hs b/dataframe-lazy/src/DataFrame/Lazy/Internal/Optimizer.hs index f8b9363a..542a2321 100644 --- a/dataframe-lazy/src/DataFrame/Lazy/Internal/Optimizer.hs +++ b/dataframe-lazy/src/DataFrame/Lazy/Internal/Optimizer.hs @@ -9,6 +9,7 @@ import qualified Data.Text as T import qualified DataFrame.Internal.Expression as E import DataFrame.Lazy.Internal.LogicalPlan import DataFrame.Lazy.Internal.PhysicalPlan +import DataFrame.Operations.Permutation (sortOrderColumns) import DataFrame.Schema (Schema (..), elements) -- | Optimise a logical plan and lower it to a physical plan. @@ -26,8 +27,8 @@ fuseFilters (Filter p1 (Filter p2 child)) = fuseFilters (Filter p child) = Filter p (fuseFilters child) fuseFilters (Project cols child) = Project cols (fuseFilters child) fuseFilters (Derive name expr child) = Derive name expr (fuseFilters child) -fuseFilters (Join jt l r left right) = - Join jt l r (fuseFilters left) (fuseFilters right) +fuseFilters (Join jt ks left right) = + Join jt ks (fuseFilters left) (fuseFilters right) fuseFilters (Aggregate keys aggs child) = Aggregate keys aggs (fuseFilters child) fuseFilters (Sort cols child) = Sort cols (fuseFilters child) @@ -62,8 +63,8 @@ pushPredicates (Filter p (Project cols child)) pushPredicates (Filter p child) = Filter p (pushPredicates child) pushPredicates (Project cols child) = Project cols (pushPredicates child) pushPredicates (Derive name expr child) = Derive name expr (pushPredicates child) -pushPredicates (Join jt l r left right) = - Join jt l r (pushPredicates left) (pushPredicates right) +pushPredicates (Join jt ks left right) = + Join jt ks (pushPredicates left) (pushPredicates right) pushPredicates (Aggregate keys aggs child) = Aggregate keys aggs (pushPredicates child) pushPredicates (Sort cols child) = Sort cols (pushPredicates child) @@ -78,16 +79,18 @@ referencedCols (Filter p child) = fmap (S.union (S.fromList (E.getColumns p))) (referencedCols child) referencedCols (Derive _ expr child) = fmap (S.union (S.fromList (someExprCols expr))) (referencedCols child) -referencedCols (Join _ l r left right) = - let keySet = S.fromList [l, r] +referencedCols (Join _ ks left right) = + let keySet = S.fromList (map fst ks <> map snd ks) lRef = fmap (S.union keySet) (referencedCols left) rRef = fmap (S.union keySet) (referencedCols right) in S.union <$> lRef <*> rRef referencedCols (Aggregate keys aggs child) = let aggCols = S.fromList (keys <> concatMap (someExprCols . snd) aggs) in fmap (S.union aggCols) (referencedCols child) -referencedCols (Sort cols child) = - fmap (S.union (S.fromList (fmap fst cols))) (referencedCols child) +referencedCols (Sort ords child) = + fmap + (S.union (S.fromList (concatMap sortOrderColumns ords))) + (referencedCols child) referencedCols (Limit _ child) = referencedCols child referencedCols (SourceDF _) = Nothing @@ -112,15 +115,18 @@ eliminateDeadColumns plan = go (referencedCols plan) plan Filter p (go (fmap (S.union (S.fromList (E.getColumns p))) needed) child) go _needed (Project cols child) = Project cols (go (Just (S.fromList cols)) child) - go needed (Join jt l r left right) = - let keySet = fmap (S.union (S.fromList [l, r])) needed - in Join jt l r (go keySet left) (go keySet right) + go needed (Join jt ks left right) = + let leftNeeded = fmap (S.union (S.fromList (map fst ks))) needed + rightNeeded = fmap (S.union (S.fromList (map snd ks))) needed + in Join jt ks (go leftNeeded left) (go rightNeeded right) go needed (Aggregate keys aggs child) = let aggCols = fmap (S.union (S.fromList (keys <> concatMap (someExprCols . snd) aggs))) needed in Aggregate keys aggs (go aggCols child) - go needed (Sort cols child) = - Sort cols (go (fmap (S.union (S.fromList (fmap fst cols))) needed) child) + go needed (Sort ords child) = + Sort + ords + (go (fmap (S.union (S.fromList (concatMap sortOrderColumns ords))) needed) child) go needed (Limit n child) = Limit n (go needed child) go needed (Scan ds schema) = case needed of @@ -170,17 +176,12 @@ toPhysical batchSz (Filter p child) = PhysicalFilter p (toPhysical batchSz child) toPhysical batchSz (Derive name expr child) = PhysicalDerive name expr (toPhysical batchSz child) -toPhysical batchSz (Join jt l r left right) = - PhysicalHashJoin - jt - l - r - (toPhysical batchSz left) - (toPhysical batchSz right) +toPhysical batchSz (Join jt ks left right) = + PhysicalHashJoin jt ks (toPhysical batchSz left) (toPhysical batchSz right) toPhysical batchSz (Aggregate keys aggs child) = PhysicalHashAggregate keys aggs (toPhysical batchSz child) -toPhysical batchSz (Sort cols child) = - PhysicalSort cols (toPhysical batchSz child) +toPhysical batchSz (Sort ords child) = + PhysicalSort ords (toPhysical batchSz child) toPhysical batchSz (Limit n child) = PhysicalLimit n (toPhysical batchSz child) toPhysical batchSz (SourceDF df) = PhysicalSourceDF batchSz df diff --git a/dataframe-lazy/src/DataFrame/Lazy/Internal/PhysicalPlan.hs b/dataframe-lazy/src/DataFrame/Lazy/Internal/PhysicalPlan.hs index e7721f70..6fdb7f02 100644 --- a/dataframe-lazy/src/DataFrame/Lazy/Internal/PhysicalPlan.hs +++ b/dataframe-lazy/src/DataFrame/Lazy/Internal/PhysicalPlan.hs @@ -3,8 +3,9 @@ module DataFrame.Lazy.Internal.PhysicalPlan where import qualified Data.Text as T import qualified DataFrame.Internal.DataFrame as D import qualified DataFrame.Internal.Expression as E -import DataFrame.Lazy.Internal.LogicalPlan (DataSource, SortOrder) +import DataFrame.Lazy.Internal.LogicalPlan (DataSource) import DataFrame.Operations.Join (JoinType) +import DataFrame.Operations.Permutation (SortOrder) import DataFrame.Schema (Schema) -- | Scan-level configuration: batch size, separator, optional pushdowns. @@ -22,10 +23,10 @@ data PhysicalPlan | PhysicalProject [T.Text] PhysicalPlan | PhysicalFilter (E.Expr Bool) PhysicalPlan | PhysicalDerive T.Text E.SomeExpr PhysicalPlan - | PhysicalHashJoin JoinType T.Text T.Text PhysicalPlan PhysicalPlan - | PhysicalSortMergeJoin JoinType T.Text T.Text PhysicalPlan PhysicalPlan + | PhysicalHashJoin JoinType [(T.Text, T.Text)] PhysicalPlan PhysicalPlan + | PhysicalSortMergeJoin JoinType [(T.Text, T.Text)] PhysicalPlan PhysicalPlan | PhysicalHashAggregate [T.Text] [(T.Text, E.SomeExpr)] PhysicalPlan - | PhysicalSort [(T.Text, SortOrder)] PhysicalPlan + | PhysicalSort [SortOrder] PhysicalPlan | PhysicalLimit Int PhysicalPlan | -- | Materialize child to a binary file on disk (used for build sides). PhysicalSpill PhysicalPlan FilePath diff --git a/dataframe-lazy/src/DataFrame/Typed/Lazy.hs b/dataframe-lazy/src/DataFrame/Typed/Lazy.hs index 7a5e19c8..03d96ff6 100644 --- a/dataframe-lazy/src/DataFrame/Typed/Lazy.hs +++ b/dataframe-lazy/src/DataFrame/Typed/Lazy.hs @@ -6,21 +6,25 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} --- | Type-safe lazy query pipelines with compile-time schema tracking. +-- | Typed lazy queries with the same operations as "DataFrame.Typed". module DataFrame.Typed.Lazy ( - -- * Core type + -- * Core types TypedLazyDataFrame, + TypedLazyGrouped, + TypedDataFrame, -- * Data sources scanCsv, + scanCsvWith, scanSeparated, + scanSeparatedWith, scanParquet, fromDataFrame, - fromTypedDataFrame, -- * Schema-preserving operations filter, take, + sortBy, -- * Schema-modifying operations derive, @@ -28,24 +32,58 @@ module DataFrame.Typed.Lazy ( -- * Aggregation groupBy, + as, aggregate, -- * Joins + JoinType (..), + join, + joinOn, innerJoin, leftJoin, rightJoin, fullOuterJoin, - -- * Sort - sortBy, - -- * Execution run, - -- * Re-exports for pipeline construction - module DataFrame.Typed.Expr, - module DataFrame.Typed.Types, - SortOrder (..), + -- * Typed expressions (combinators live in "DataFrame.Typed.Functions") + TExpr (..), + TSortOrder (..), + col, + lit, + ifThenElse, + caseWhen, + orElse, + (.==.), + (./=.), + (.<.), + (.<=.), + (.>=.), + (.>.), + (.+.), + (.-.), + (.*.), + (./.), + (.^^.), + (.^.), + (.+), + (.-), + (.*), + (./), + (.^^), + (.^), + (.==), + (./=), + (.<), + (.<=), + (.>=), + (.>), + (.&&.), + (.||.), + (.&&), + (.||), + (|>), ) where import Data.Kind (Type) @@ -54,15 +92,25 @@ import qualified Data.Text as T import GHC.TypeLits (KnownSymbol, Symbol, symbolVal) import Prelude hiding (filter, take) +import DataFrame.Expression.Operators ((|>)) +import DataFrame.IO.CSV (CsvReader) import qualified DataFrame.Internal.Column as C -import qualified DataFrame.Internal.Expression as E import DataFrame.Lazy.Internal.DataFrame (LazyDataFrame) import qualified DataFrame.Lazy.Internal.DataFrame as L -import DataFrame.Lazy.Internal.LogicalPlan (SortOrder (..)) import DataFrame.Operations.Join (JoinType (..)) -import DataFrame.Schema (Schema) +import DataFrame.Schema (RuntimeSchema (..)) +import DataFrame.Typed.Aggregate (as) import DataFrame.Typed.Expr import DataFrame.Typed.Freeze (unsafeFreeze) +import DataFrame.Typed.Join ( + Fsts, + JoinKeys, + JoinOnSchema, + JoinSchema, + KnownJoinType (..), + SwapPairs, + ) +import DataFrame.Typed.Operations (AllKnownPairs (..), toUntypedSort) import DataFrame.Typed.Schema import DataFrame.Typed.Types @@ -72,36 +120,39 @@ newtype TypedLazyDataFrame (cols :: [(Symbol, Type)]) = TLD {_unTLD :: LazyDataF instance Show (TypedLazyDataFrame cols) where show (TLD ldf) = "TypedLazyDataFrame { " ++ show ldf ++ " }" --- | Scan a CSV file with a given schema. +-- | Scan a CSV file; the column types come from @cols@. scanCsv :: - Schema -> - T.Text -> - TypedLazyDataFrame cols -scanCsv schema path = TLD (L.scanCsv schema path) + forall cols. (RuntimeSchema cols) => T.Text -> TypedLazyDataFrame cols +scanCsv path = TLD (L.scanCsv (runtimeSchema @cols) path) + +-- | Like 'scanCsv' but with an explicit CSV reader. +scanCsvWith :: + forall cols. + (RuntimeSchema cols) => CsvReader -> T.Text -> TypedLazyDataFrame cols +scanCsvWith reader path = TLD (L.scanCsvWith reader (runtimeSchema @cols) path) --- | Scan a character-separated file with a given schema. +-- | Scan a character-separated file; the column types come from @cols@. scanSeparated :: - Char -> - Schema -> - T.Text -> - TypedLazyDataFrame cols -scanSeparated sep schema path = TLD (L.scanSeparated sep schema path) + forall cols. (RuntimeSchema cols) => Char -> T.Text -> TypedLazyDataFrame cols +scanSeparated sep path = TLD (L.scanSeparated sep (runtimeSchema @cols) path) + +-- | Like 'scanSeparated' but with an explicit CSV reader. +scanSeparatedWith :: + forall cols. + (RuntimeSchema cols) => + CsvReader -> Char -> T.Text -> TypedLazyDataFrame cols +scanSeparatedWith reader sep path = + TLD (L.scanSeparatedWith reader sep (runtimeSchema @cols) path) --- | Scan a Parquet file, directory, or glob pattern with a given schema. +-- | Scan a Parquet file, directory, or glob; the columns come from @cols@. scanParquet :: - Schema -> - T.Text -> - TypedLazyDataFrame cols -scanParquet schema path = TLD (L.scanParquet schema path) + forall cols. (RuntimeSchema cols) => T.Text -> TypedLazyDataFrame cols +scanParquet path = TLD (L.scanParquet (runtimeSchema @cols) path) -- | Lift an already-loaded eager 'TypedDataFrame' into a lazy plan. fromDataFrame :: TypedDataFrame cols -> TypedLazyDataFrame cols fromDataFrame (TDF df) = TLD (L.fromDataFrame df) --- | Synonym for 'fromDataFrame'. -fromTypedDataFrame :: TypedDataFrame cols -> TypedLazyDataFrame cols -fromTypedDataFrame = fromDataFrame - -- | Keep rows that satisfy the predicate. filter :: TExpr cols Bool -> TypedLazyDataFrame cols -> TypedLazyDataFrame cols filter (TExpr expr) (TLD ldf) = TLD (L.filter expr ldf) @@ -110,6 +161,11 @@ filter (TExpr expr) (TLD ldf) = TLD (L.filter expr ldf) take :: Int -> TypedLazyDataFrame cols -> TypedLazyDataFrame cols take n (TLD ldf) = TLD (L.take n ldf) +-- | Sort by the given typed orders, as the eager typed @sortBy@ does. +sortBy :: + [TSortOrder cols] -> TypedLazyDataFrame cols -> TypedLazyDataFrame cols +sortBy ords (TLD ldf) = TLD (L.sortBy (map toUntypedSort ords) ldf) + -- | Add a computed column. derive :: forall name a cols. @@ -126,112 +182,86 @@ select :: (AllKnownSymbol names, AssertAllPresent names cols) => TypedLazyDataFrame cols -> TypedLazyDataFrame (SubsetSchema names cols) -select (TLD ldf) = TLD (L.select (DataFrame.Typed.Schema.symbolVals @names) ldf) +select (TLD ldf) = TLD (L.select (symbolVals @names) ldf) --- | A typed lazy grouped query. -newtype TypedLazyGrouped (keys :: [Symbol]) (cols :: [(Symbol, Type)]) = TLG - { _unTLG :: ([T.Text], LazyDataFrame) - } +-- | A typed lazy query grouped by @keys@, awaiting 'aggregate'. +newtype TypedLazyGrouped (keys :: [Symbol]) (cols :: [(Symbol, Type)]) + = TLG L.LazyGroupedDataFrame --- | Group by key columns. +-- | Group by key columns; finish with 'aggregate'. groupBy :: forall (keys :: [Symbol]) cols. (AllKnownSymbol keys, AssertAllPresent keys cols) => TypedLazyDataFrame cols -> TypedLazyGrouped keys cols -groupBy (TLD ldf) = TLG (DataFrame.Typed.Schema.symbolVals @keys, ldf) +groupBy (TLD ldf) = TLG (L.groupBy (symbolVals @keys) ldf) --- | Aggregate a grouped lazy query. +-- | Aggregate each group; build the outputs with 'as', as in the eager API. aggregate :: forall keys cols aggs. - TAgg keys cols aggs -> + (TAgg keys cols '[] -> TAgg keys cols aggs) -> TypedLazyGrouped keys cols -> TypedLazyDataFrame (Append (GroupKeyColumns keys cols) (Reverse aggs)) -aggregate tagg (TLG (keys, ldf)) = - TLD (L.groupBy keys (aggToNamedExprs tagg) ldf) +aggregate build (TLG grouped) = + TLD (L.aggregate (taggToNamedExprs (build TAggNil)) grouped) --- | Typed inner join on a single key column present in both schemas. -innerJoin :: - forall (key :: Symbol) left right. - ( KnownSymbol key - , AssertAllPresent '[key] left - , AssertAllPresent '[key] right - , AssertKeyTypesMatch '[key] left right - ) => +-- | Typed lazy join: @join \@'INNER \@'["key"] l r@. +join :: + forall (jt :: JoinType) (keys :: [Symbol]) left right. + (KnownJoinType jt, JoinKeys keys left right) => TypedLazyDataFrame left -> TypedLazyDataFrame right -> - TypedLazyDataFrame (InnerJoinSchema '[key] left right) -innerJoin = joinOn @key INNER + TypedLazyDataFrame (JoinSchema jt keys left right) +join (TLD l) (TLD r) = TLD (L.join (joinTypeVal @jt) (symbolVals @keys) l r) --- | Typed left join. The right table's non-key columns become @Maybe@. -leftJoin :: - forall (key :: Symbol) left right. - ( KnownSymbol key - , AssertAllPresent '[key] left - , AssertAllPresent '[key] right - , AssertKeyTypesMatch '[key] left right +-- | Join on @(leftKey, rightKey)@ pairs; the result keeps the left names. +joinOn :: + forall (jt :: JoinType) (pairs :: [(Symbol, Symbol)]) left right. + ( KnownJoinType jt + , AllKnownPairs pairs + , JoinKeys (Fsts pairs) left (RenameManyInSchema (SwapPairs pairs) right) ) => TypedLazyDataFrame left -> TypedLazyDataFrame right -> - TypedLazyDataFrame (LeftJoinSchema '[key] left right) -leftJoin = joinOn @key LEFT + TypedLazyDataFrame (JoinOnSchema jt pairs left right) +joinOn (TLD l) (TLD r) = TLD (L.joinOn (joinTypeVal @jt) (pairVals @pairs) l r) --- | Typed right join. The left table's non-key columns become @Maybe@. -rightJoin :: - forall (key :: Symbol) left right. - ( KnownSymbol key - , AssertAllPresent '[key] left - , AssertAllPresent '[key] right - , AssertKeyTypesMatch '[key] left right - ) => +-- | Typed inner join on one or more key columns. +innerJoin :: + forall (keys :: [Symbol]) left right. + (JoinKeys keys left right) => TypedLazyDataFrame left -> TypedLazyDataFrame right -> - TypedLazyDataFrame (RightJoinSchema '[key] left right) -rightJoin = joinOn @key RIGHT + TypedLazyDataFrame (InnerJoinSchema keys left right) +innerJoin = join @'INNER @keys --- | Typed full outer join. Non-key columns from both tables become @Maybe@. -fullOuterJoin :: - forall (key :: Symbol) left right. - ( KnownSymbol key - , AssertAllPresent '[key] left - , AssertAllPresent '[key] right - , AssertKeyTypesMatch '[key] left right - ) => +-- | Typed left join; the right side's non-key columns become @Maybe@. +leftJoin :: + forall (keys :: [Symbol]) left right. + (JoinKeys keys left right) => TypedLazyDataFrame left -> TypedLazyDataFrame right -> - TypedLazyDataFrame (FullOuterJoinSchema '[key] left right) -fullOuterJoin = joinOn @key FULL_OUTER + TypedLazyDataFrame (LeftJoinSchema keys left right) +leftJoin = join @'LEFT @keys --- | Runtime delegation shared by the typed joins. -joinOn :: - forall (key :: Symbol) left right out. - (KnownSymbol key) => - JoinType -> +-- | Typed right join; the left side's non-key columns become @Maybe@. +rightJoin :: + forall (keys :: [Symbol]) left right. + (JoinKeys keys left right) => TypedLazyDataFrame left -> TypedLazyDataFrame right -> - TypedLazyDataFrame out -joinOn jt (TLD left) (TLD right) = - TLD (L.join jt keyName keyName left right) - where - keyName = T.pack (symbolVal (Proxy @key)) + TypedLazyDataFrame (RightJoinSchema keys left right) +rightJoin = join @'RIGHT @keys --- | Sort the result by column name and direction. -sortBy :: - [(T.Text, SortOrder)] -> - TypedLazyDataFrame cols -> - TypedLazyDataFrame cols -sortBy cols (TLD ldf) = TLD (L.sortBy cols ldf) +-- | Typed full outer join; non-key columns on both sides become @Maybe@. +fullOuterJoin :: + forall (keys :: [Symbol]) left right. + (JoinKeys keys left right) => + TypedLazyDataFrame left -> + TypedLazyDataFrame right -> + TypedLazyDataFrame (FullOuterJoinSchema keys left right) +fullOuterJoin = join @'FULL_OUTER @keys -- | Execute the lazy query and return a typed DataFrame. -run :: - forall cols. - (KnownSchema cols) => - TypedLazyDataFrame cols -> - IO (TypedDataFrame cols) -run (TLD ldf) = unsafeFreeze <$> L.runDataFrame ldf - --- | Convert TAgg to untyped named expressions for the lazy groupBy. -aggToNamedExprs :: TAgg keys cols aggs -> [(T.Text, E.SomeExpr)] -aggToNamedExprs TAggNil = [] -aggToNamedExprs (TAggCons name (TExpr expr) rest) = - (name, E.SomeExpr expr) : aggToNamedExprs rest +run :: TypedLazyDataFrame cols -> IO (TypedDataFrame cols) +run (TLD ldf) = unsafeFreeze <$> L.run ldf diff --git a/dataframe-lazy/test/StreamingCsv.hs b/dataframe-lazy/test/StreamingCsv.hs index 3031263d..74a22093 100644 --- a/dataframe-lazy/test/StreamingCsv.hs +++ b/dataframe-lazy/test/StreamingCsv.hs @@ -46,7 +46,7 @@ streamingPathReaderFillsWindow = modifyIORef' calls (+ 1) Csv.readSeparated opts windowPath df <- - L.runDataFrame + L.run (L.scanCsvStreamingWith countingReader schema (T.pack path)) assertEqual "all rows decoded" (rowCount, 1) (D.dataframeDimensions df) readIORef calls >>= assertEqual "one parser window" 1 @@ -73,7 +73,7 @@ streamingWindowBoundaryIsLossless = modifyIORef' observations (observation :) pure D.empty _ <- - L.runDataFrame + L.run (L.scanCsvStreamingBytesWith observingReader schema (T.pack path)) actual <- reverse <$> readIORef observations assertEqual @@ -98,7 +98,7 @@ streamingQuotedNewlineSpansWindow = modifyIORef' bodyLengths (BS.length body :) pure D.empty _ <- - L.runDataFrame + L.run (L.scanCsvStreamingBytesWith observingReader schema (T.pack path)) actual <- reverse <$> readIORef bodyLengths assertEqual @@ -121,7 +121,7 @@ streamingDefaultMatchesEager = withRawCsv input $ \path -> do let opts = Csv.schemaReadOptions schema eager <- Csv.readSeparated opts path - streamed <- L.runDataFrame (L.scanCsvStreaming schema (T.pack path)) + streamed <- L.run (L.scanCsvStreaming schema (T.pack path)) assertEqual "streaming bytes reader matches eager reader" eager streamed -- | A parser failure on the producer thread must fail the consumer. @@ -134,7 +134,7 @@ streamingReaderFailurePropagates = withRawCsv input $ \path -> do result <- try - ( L.runDataFrame + ( L.run (L.scanCsvStreamingBytesWith failingReader schema (T.pack path)) ) case result of @@ -155,7 +155,7 @@ streamingKeepsTextPacked = schema = Schema $ M.fromList [("station", schemaType @T.Text)] withRawCsv input $ \path -> do df <- - L.runDataFrame + L.run (L.scanCsvStreaming schema (T.pack path)) case D.unsafeGetColumn "station" df of PackedText{} -> pure () diff --git a/dataframe-learn/README.md b/dataframe-learn/README.md index 043a1eff..9d714f60 100644 --- a/dataframe-learn/README.md +++ b/dataframe-learn/README.md @@ -198,7 +198,6 @@ D.columnNames (D.derive "prediction" (predict model) sales) > > :87:47: error: [GHC-88464] > Variable not in scope: model -> Suggested fix: Perhaps use ‘T.mode’ (imported from DataFrame.Typed) > > :87:54: error: [GHC-88464] > Variable not in scope: sales :: D.DataFrame diff --git a/dataframe-learn/src-internal/DataFrame/DecisionTree/Fit.hs b/dataframe-learn/src-internal/DataFrame/DecisionTree/Fit.hs index 103cc003..076b9232 100644 --- a/dataframe-learn/src-internal/DataFrame/DecisionTree/Fit.hs +++ b/dataframe-learn/src-internal/DataFrame/DecisionTree/Fit.hs @@ -44,7 +44,8 @@ import DataFrame.Internal.DataFrame (DataFrame) import DataFrame.Internal.Expression (Expr (..)) import DataFrame.Internal.Interpreter (interpret) import DataFrame.Operations.Core (nRows) -import DataFrame.Operations.Subset (exclude, filterWhere) +import DataFrame.Operations.Subset (exclude) +import qualified DataFrame.Operations.Subset as Subset import Control.Exception (throw) import Data.Function (on) @@ -101,7 +102,7 @@ pruneTree :: forall a. (Columnable a) => Expr a -> Expr a pruneTree = pruneExpr partitionDataFrame :: Expr Bool -> DataFrame -> (DataFrame, DataFrame) -partitionDataFrame cond df = (filterWhere cond df, filterWhere (F.not cond) df) +partitionDataFrame cond df = (Subset.filter cond df, Subset.filter (F.not cond) df) -- | Laplace-smoothed Gini impurity of the target distribution. calculateGini :: diff --git a/dataframe-operations/dataframe-operations.cabal b/dataframe-operations/dataframe-operations.cabal index 2215163e..245bc8a8 100644 --- a/dataframe-operations/dataframe-operations.cabal +++ b/dataframe-operations/dataframe-operations.cabal @@ -52,6 +52,7 @@ library DataFrame.Typed.Apply DataFrame.Typed.Expr DataFrame.Typed.Expr.Extra + DataFrame.Typed.Functions DataFrame.Typed.Join DataFrame.Typed.Operations DataFrame.Typed.Sampling diff --git a/dataframe-operations/src/DataFrame/Monad.hs b/dataframe-operations/src/DataFrame/Monad.hs index 49d63b95..e298e375 100644 --- a/dataframe-operations/src/DataFrame/Monad.hs +++ b/dataframe-operations/src/DataFrame/Monad.hs @@ -20,7 +20,7 @@ module DataFrame.Monad ( deriveM, insertM, renameM, - filterWhereM, + filterM, sampleM, takeM, dropM, @@ -103,8 +103,8 @@ renameM :: (Columnable a) => Expr a -> T.Text -> FrameM (Expr a) renameM (Col oldName) newName = execWithExpr newName (D.rename oldName newName) renameM expr newName = deriveM newName expr -filterWhereM :: Expr Bool -> FrameM () -filterWhereM p = modifyM (D.filterWhere p) +filterM :: Expr Bool -> FrameM () +filterM p = modifyM (D.filter p) sampleM :: (RandomGen g) => g -> Double -> FrameM () sampleM pureGen p = modifyM (D.sample pureGen p) diff --git a/dataframe-operations/src/DataFrame/Operations/Join.hs b/dataframe-operations/src/DataFrame/Operations/Join.hs index 1acb907f..ce10fed7 100644 --- a/dataframe-operations/src/DataFrame/Operations/Join.hs +++ b/dataframe-operations/src/DataFrame/Operations/Join.hs @@ -14,6 +14,7 @@ module DataFrame.Operations.Join ( -- * Joins join, + joinOn, innerJoin, leftJoin, rightJoin, @@ -99,6 +100,20 @@ join LEFT xs right = leftJoin xs right join RIGHT xs right = rightJoin xs right join FULL_OUTER xs right = fullOuterJoin xs right +-- | Join on @(leftKey, rightKey)@ pairs; the result keeps the left names. +joinOn :: + JoinType -> + [(T.Text, T.Text)] -> + DataFrame -> + DataFrame -> + DataFrame +joinOn jt keyPairs left right = + join + jt + (map fst keyPairs) + left + (renameMany [(r, l) | (l, r) <- keyPairs, l /= r] right) + -- | Build-side row count above which the hash join switches strategy. joinStrategyThreshold :: Int joinStrategyThreshold = 500_000 diff --git a/dataframe-operations/src/DataFrame/Operations/Permutation.hs b/dataframe-operations/src/DataFrame/Operations/Permutation.hs index d4a81371..c2f086f3 100644 --- a/dataframe-operations/src/DataFrame/Operations/Permutation.hs +++ b/dataframe-operations/src/DataFrame/Operations/Permutation.hs @@ -9,6 +9,7 @@ module DataFrame.Operations.Permutation ( -- * Sorting SortOrder (..), + sortOrderColumns, sortBy, topBy, @@ -75,6 +76,12 @@ instance Eq SortOrder where (==) (Desc _) (Desc _) = True (==) _ _ = False +instance Show SortOrder where + showsPrec d (Asc e) = showParen (d > 10) (showString "Asc " . showsPrec 11 e) + showsPrec d (Desc e) = + showParen (d > 10) (showString "Desc " . showsPrec 11 e) + +-- | Columns a sort order reads. sortOrderColumns :: SortOrder -> [T.Text] sortOrderColumns (Asc e) = getColumns e sortOrderColumns (Desc e) = getColumns e diff --git a/dataframe-operations/src/DataFrame/Operations/Subset.hs b/dataframe-operations/src/DataFrame/Operations/Subset.hs index 4ab15b22..204474ee 100644 --- a/dataframe-operations/src/DataFrame/Operations/Subset.hs +++ b/dataframe-operations/src/DataFrame/Operations/Subset.hs @@ -23,7 +23,6 @@ module DataFrame.Operations.Subset ( -- * Filtering filter, filterBy, - filterWhere, filterJust, filterNothing, filterAllJust, @@ -190,25 +189,28 @@ range (start, end) d = clip :: Int -> Int -> Int -> Int clip n left right = min right $ max n left --- | O(n * k) Filter rows by a given condition. -filter :: +-- | O(n * k) Keep rows whose value of the expression satisfies the predicate. +filterBy :: forall a. (Columnable a) => - -- | Column to filter by - Expr a -> -- | Filter condition (a -> Bool) -> + -- | Column to filter by + Expr a -> -- | Dataframe to filter DataFrame -> DataFrame -filter e@(Col filterColumnName) condition df = case getColumn filterColumnName df of +filterBy condition e@(Col filterColumnName) df = case getColumn filterColumnName df of Nothing -> throw $ - ColumnsNotFoundException [filterColumnName] "filter" (M.keys $ columnIndices df) + ColumnsNotFoundException + [filterColumnName] + "filterBy" + (M.keys $ columnIndices df) Just c@(PackedText _ _) -> - filter e condition (insertColumn filterColumnName (materializePacked c) df) + filterBy condition e (insertColumn filterColumnName (materializePacked c) df) Just c@(MergedColumn _ _) -> - filter e condition (insertColumn filterColumnName (materializeMerged c) df) + filterBy condition e (insertColumn filterColumnName (materializeMerged c) df) Just _col@(BoxedColumn bm (column :: V.Vector b)) -> case testEquality (typeRep @a) (typeRep @b) of Just Refl -> filterByVector filterColumnName column condition df @@ -234,7 +236,7 @@ filter e@(Col filterColumnName) condition df = case getColumn filterColumnName d Nothing -> filterByVector filterColumnName column condition df Nothing -> filterByVector filterColumnName column condition df _ -> filterByVector filterColumnName column condition df -filter expr condition df = +filterBy condition expr df = let (TColumn col') = case interpret @a df (normalize expr) of Left e -> throw e @@ -261,7 +263,7 @@ filterByVector filterColumnName column condition df = case testEquality (typeRep { userType = Right $ typeRep @a , expectedType = Right $ typeRep @b , errorColumnName = Just (T.unpack filterColumnName) - , callingFunctionName = Just "filter" + , callingFunctionName = Just "filterBy" } ) Just Refl -> @@ -273,13 +275,9 @@ filterByVector filterColumnName column condition df = case testEquality (typeRep , dataframeDimensions = (VG.length ixs, snd (dataframeDimensions df)) } --- | O(k) A version of filter where the predicate comes first. -filterBy :: (Columnable a) => (a -> Bool) -> Expr a -> DataFrame -> DataFrame -filterBy = flip filter - -- | O(k) Filter the dataframe with a boolean expression. -filterWhere :: Expr Bool -> DataFrame -> DataFrame -filterWhere expr df = +filter :: Expr Bool -> DataFrame -> DataFrame +filter expr df = let (TColumn col') = case interpret @Bool df (normalize expr) of Left e -> throw e @@ -305,9 +303,9 @@ filterJust colName df = case getColumn colName df of 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 + filterBy isJust (Col @(Maybe a) colName) df & apply @(Maybe a) fromJust colName UnboxedColumn (Just _) (_col :: VU.Vector a) -> - filter (Col @(Maybe a) colName) isJust df & apply @(Maybe a) fromJust colName + filterBy isJust (Col @(Maybe a) colName) df & apply @(Maybe a) fromJust colName c@(PackedText (Just _) _) -> filterJust colName (insertColumn colName (materializePacked c) df) _ -> df @@ -320,8 +318,8 @@ filterNothing colName df = case getColumn colName df of throw $ ColumnsNotFoundException [colName] "filterNothing" (M.keys $ columnIndices df) Just column | hasMissing column -> case column of - BoxedColumn (Just _) (_col :: V.Vector a) -> filter (Col @(Maybe a) colName) isNothing df - UnboxedColumn (Just _) (_col :: VU.Vector a) -> filter (Col @(Maybe a) colName) isNothing df + BoxedColumn (Just _) (_col :: V.Vector a) -> filterBy isNothing (Col @(Maybe a) colName) df + UnboxedColumn (Just _) (_col :: VU.Vector a) -> filterBy isNothing (Col @(Maybe a) colName) df c@(PackedText (Just _) _) -> filterNothing colName (insertColumn colName (materializePacked c) df) _ -> df @@ -435,7 +433,7 @@ sample pureGen p df = in df & insertColumn (name cRand) rand - & filterWhere (cRand .>=. Lit (1 - p)) + & filter (cRand .>=. Lit (1 - p)) & exclude [name cRand] -- | Split a dataset in two: a sample of fraction p and the remainder. @@ -448,10 +446,10 @@ randomSplit pureGen p df = withRand = df & insertColumn (name cRand) rand in ( withRand - & filterWhere (cRand .<=. Lit p) + & filter (cRand .<=. Lit p) & exclude [name cRand] , withRand - & filterWhere + & filter (cRand .>. Lit p) & exclude [name cRand] ) @@ -465,12 +463,12 @@ kFolds pureGen folds df = withRand = df & insertColumn (name cRand) rand partitionSize = 1 / fromIntegral folds singleFold n d = - d & filterWhere (cRand .>=. Lit (fromIntegral n * partitionSize)) + d & filter (cRand .>=. Lit (fromIntegral n * partitionSize)) go (-1) _ = [] go n d = let d' = singleFold n d - d'' = d & filterWhere (cRand .<. Lit (fromIntegral n * partitionSize)) + d'' = d & filter (cRand .<. Lit (fromIntegral n * partitionSize)) in d' : go (n - 1) d'' in diff --git a/dataframe-operations/src/DataFrame/Typed/Apply.hs b/dataframe-operations/src/DataFrame/Typed/Apply.hs index bfb4cda8..d1e0ad88 100644 --- a/dataframe-operations/src/DataFrame/Typed/Apply.hs +++ b/dataframe-operations/src/DataFrame/Typed/Apply.hs @@ -8,7 +8,9 @@ -- | Typed column transformations that track schema changes in the types. module DataFrame.Typed.Apply ( - applyColumn, + apply, + applyInt, + applyDouble, applyMany, applyWhere, applyAtIndex, @@ -47,7 +49,7 @@ import DataFrame.Typed.Schema ( import DataFrame.Typed.Types (TExpr (..), TypedDataFrame (..)) -- | Map a function over a column, rewriting its element type from @a@ to @b@. -applyColumn :: +apply :: forall name a b cols. ( KnownSymbol name , a ~ SafeLookup name cols @@ -58,11 +60,37 @@ applyColumn :: (a -> b) -> TypedDataFrame cols -> TypedDataFrame (SetColumnType name b cols) -applyColumn f (TDF df) = unsafeFreeze (D.apply f colName df) +apply f (TDF df) = unsafeFreeze (D.apply f colName df) where colName = T.pack (symbolVal (Proxy @name)) --- | Like 'applyColumn' but returns the error instead of throwing. +-- | 'apply' specialised to an @Int@ column. +applyInt :: + forall name b cols. + ( KnownSymbol name + , Int ~ SafeLookup name cols + , Columnable b + , AssertPresent name cols + ) => + (Int -> b) -> + TypedDataFrame cols -> + TypedDataFrame (SetColumnType name b cols) +applyInt = apply @name + +-- | 'apply' specialised to a @Double@ column. +applyDouble :: + forall name b cols. + ( KnownSymbol name + , Double ~ SafeLookup name cols + , Columnable b + , AssertPresent name cols + ) => + (Double -> b) -> + TypedDataFrame cols -> + TypedDataFrame (SetColumnType name b cols) +applyDouble = apply @name + +-- | Like 'apply' but returns the error instead of throwing. safeApply :: forall name a b cols. ( KnownSymbol name diff --git a/dataframe-operations/src/DataFrame/Typed/Functions.hs b/dataframe-operations/src/DataFrame/Typed/Functions.hs new file mode 100644 index 00000000..20385997 --- /dev/null +++ b/dataframe-operations/src/DataFrame/Typed/Functions.hs @@ -0,0 +1,6 @@ +-- | Typed expression combinators; the typed "DataFrame.Functions". +module DataFrame.Typed.Functions ( + module DataFrame.Typed.Expr, +) where + +import DataFrame.Typed.Expr diff --git a/dataframe-operations/src/DataFrame/Typed/Join.hs b/dataframe-operations/src/DataFrame/Typed/Join.hs index ffb6c4d2..cec60e56 100644 --- a/dataframe-operations/src/DataFrame/Typed/Join.hs +++ b/dataframe-operations/src/DataFrame/Typed/Join.hs @@ -1,25 +1,39 @@ {-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE DataKinds #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE GADTs #-} +{-# LANGUAGE KindSignatures #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} {-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} module DataFrame.Typed.Join ( -- * Typed joins + join, + joinOn, + JoinSchema, + JoinOnSchema, + JoinKeys, + KnownJoinType (..), + Fsts, + SwapPairs, innerJoin, leftJoin, rightJoin, fullOuterJoin, ) where +import Data.Kind (Type) import GHC.TypeLits (Symbol) +import DataFrame.Operations.Join (JoinType (..)) import qualified DataFrame.Operations.Join as DJ import DataFrame.Typed.Freeze (unsafeFreeze) +import DataFrame.Typed.Operations (AllKnownPairs (..)) import DataFrame.Typed.Schema import DataFrame.Typed.Types (TypedDataFrame (..)) @@ -35,7 +49,7 @@ innerJoin :: TypedDataFrame right -> TypedDataFrame (InnerJoinSchema keys left right) innerJoin (TDF l) (TDF r) = - unsafeFreeze (DJ.innerJoin keyNames r l) + unsafeFreeze (DJ.innerJoin keyNames l r) where keyNames = symbolVals @keys @@ -83,6 +97,74 @@ fullOuterJoin :: TypedDataFrame right -> TypedDataFrame (FullOuterJoinSchema keys left right) fullOuterJoin (TDF l) (TDF r) = - unsafeFreeze (DJ.fullOuterJoin keyNames r l) + unsafeFreeze (DJ.fullOuterJoin keyNames l r) where keyNames = symbolVals @keys + +-- | Keys present on both sides with matching types. +type JoinKeys keys left right = + ( AllKnownSymbol keys + , AssertAllPresent keys left + , AssertAllPresent keys right + , AssertKeyTypesMatch keys left right + ) + +-- | Result schema of a @jt@ join on @keys@. +type family + JoinSchema + (jt :: JoinType) + (keys :: [Symbol]) + (left :: [(Symbol, Type)]) + (right :: [(Symbol, Type)]) :: + [(Symbol, Type)] + where + JoinSchema 'INNER keys left right = InnerJoinSchema keys left right + JoinSchema 'LEFT keys left right = LeftJoinSchema keys left right + JoinSchema 'RIGHT keys left right = RightJoinSchema keys left right + JoinSchema 'FULL_OUTER keys left right = FullOuterJoinSchema keys left right + +-- | Reflect a promoted 'JoinType' to its value. +class KnownJoinType (jt :: JoinType) where + joinTypeVal :: JoinType + +instance KnownJoinType 'INNER where joinTypeVal = INNER +instance KnownJoinType 'LEFT where joinTypeVal = LEFT +instance KnownJoinType 'RIGHT where joinTypeVal = RIGHT +instance KnownJoinType 'FULL_OUTER where joinTypeVal = FULL_OUTER + +-- | Typed counterpart of untyped @join@: @join \@'INNER \@'["key"] l r@. +join :: + forall (jt :: JoinType) (keys :: [Symbol]) left right. + (KnownJoinType jt, JoinKeys keys left right) => + TypedDataFrame left -> + TypedDataFrame right -> + TypedDataFrame (JoinSchema jt keys left right) +join (TDF l) (TDF r) = + unsafeFreeze (DJ.join (joinTypeVal @jt) (symbolVals @keys) l r) + +-- | Left names of @(leftKey, rightKey)@ pairs. +type family Fsts (pairs :: [(Symbol, Symbol)]) :: [Symbol] where + Fsts '[] = '[] + Fsts ('(l, _) ': rest) = l ': Fsts rest + +-- | Flip each pair, turning @(left, right)@ into a right-to-left rename. +type family SwapPairs (pairs :: [(Symbol, Symbol)]) :: [(Symbol, Symbol)] where + SwapPairs '[] = '[] + SwapPairs ('(l, r) ': rest) = '(r, l) ': SwapPairs rest + +-- | Result schema of 'joinOn': right keys take their left names first. +type JoinOnSchema jt pairs left right = + JoinSchema jt (Fsts pairs) left (RenameManyInSchema (SwapPairs pairs) right) + +-- | Join on key pairs: @joinOn \@'INNER \@'[ '("id", "uid")] l r@. +joinOn :: + forall (jt :: JoinType) (pairs :: [(Symbol, Symbol)]) left right. + ( KnownJoinType jt + , AllKnownPairs pairs + , JoinKeys (Fsts pairs) left (RenameManyInSchema (SwapPairs pairs) right) + ) => + TypedDataFrame left -> + TypedDataFrame right -> + TypedDataFrame (JoinOnSchema jt pairs left right) +joinOn (TDF l) (TDF r) = + unsafeFreeze (DJ.joinOn (joinTypeVal @jt) (pairVals @pairs) l r) diff --git a/dataframe-operations/src/DataFrame/Typed/Operations.hs b/dataframe-operations/src/DataFrame/Typed/Operations.hs index b276ea4d..8e756ffc 100644 --- a/dataframe-operations/src/DataFrame/Typed/Operations.hs +++ b/dataframe-operations/src/DataFrame/Typed/Operations.hs @@ -11,7 +11,6 @@ module DataFrame.Typed.Operations ( -- * Schema-preserving operations - filterWhere, filter, filterBy, filterAllJust, @@ -19,6 +18,7 @@ module DataFrame.Typed.Operations ( filterNothing, filterAllNothing, sortBy, + topBy, take, takeLast, drop, @@ -31,6 +31,7 @@ module DataFrame.Typed.Operations ( -- * Schema-modifying operations derive, + deriveMany, impute, select, exclude, @@ -47,6 +48,11 @@ module DataFrame.Typed.Operations ( valueCounts, valueProportions, + -- * Construction and folds + empty, + null, + fold, + -- * Metadata dimensions, nRows, @@ -61,14 +67,19 @@ module DataFrame.Typed.Operations ( intersect, difference, symmetricDifference, + + -- * Helpers + AllKnownPairs (..), + toUntypedSort, ) where +import qualified Data.List as L import Data.Proxy (Proxy (..)) import qualified Data.Text as T import qualified Data.Vector as V import GHC.TypeLits (KnownSymbol, Symbol, symbolVal) import System.Random (RandomGen) -import Prelude hiding (drop, filter, take) +import Prelude hiding (drop, filter, null, take) import qualified DataFrame.Functions as DF import DataFrame.Internal.Column (Columnable) @@ -84,19 +95,19 @@ import qualified DataFrame.Operations.Transformations as D import DataFrame.Typed.Freeze (unsafeFreeze) import DataFrame.Typed.Schema -import DataFrame.Typed.Types (TExpr (..), TSortOrder (..), TypedDataFrame (..)) +import DataFrame.Typed.Types ( + TAgg (..), + TExpr (..), + TSortOrder (..), + TypedDataFrame (..), + taggToNamedExprs, + ) -- | Filter rows where a boolean expression evaluates to True. -filterWhere :: TExpr cols Bool -> TypedDataFrame cols -> TypedDataFrame cols -filterWhere (TExpr expr) (TDF df) = TDF (D.filterWhere expr df) +filter :: TExpr cols Bool -> TypedDataFrame cols -> TypedDataFrame cols +filter (TExpr expr) (TDF df) = TDF (D.filter expr df) -- | Filter rows by applying a predicate to a typed expression. -filter :: - (Columnable a) => - TExpr cols a -> (a -> Bool) -> TypedDataFrame cols -> TypedDataFrame cols -filter (TExpr expr) pred' (TDF df) = TDF (D.filter expr pred' df) - --- | Filter rows by a predicate on a column expression (flipped argument order). filterBy :: (Columnable a) => (a -> Bool) -> TExpr cols a -> TypedDataFrame cols -> TypedDataFrame cols @@ -135,10 +146,16 @@ filterAllNothing (TDF df) = TDF (D.filterAllNothing df) -- | Sort by the given typed sort orders, validated against the schema. sortBy :: [TSortOrder cols] -> TypedDataFrame cols -> TypedDataFrame cols sortBy ords (TDF df) = TDF (D.sortBy (map toUntypedSort ords) df) - where - toUntypedSort :: TSortOrder cols -> D.SortOrder - toUntypedSort (Asc (TExpr e)) = D.Asc e - toUntypedSort (Desc (TExpr e)) = D.Desc e + +-- | Rows @[offset, offset + count)@ of the frame sorted by the given orders. +topBy :: + Int -> Int -> [TSortOrder cols] -> TypedDataFrame cols -> TypedDataFrame cols +topBy offset count ords (TDF df) = + TDF (D.topBy offset count (map toUntypedSort ords) df) + +toUntypedSort :: TSortOrder cols -> D.SortOrder +toUntypedSort (Asc (TExpr e)) = D.Asc e +toUntypedSort (Desc (TExpr e)) = D.Desc e -- | Take the first @n@ rows. take :: Int -> TypedDataFrame cols -> TypedDataFrame cols @@ -191,6 +208,16 @@ derive (TExpr expr) (TDF df) = unsafeFreeze (D.derive colName expr df) where colName = T.pack (symbolVal (Proxy @name)) +-- | Add several columns at once; build them with 'DataFrame.Typed.as'. +deriveMany :: + forall new cols. + (AssertAllAbsent (Reverse new) cols) => + (TAgg '[] cols '[] -> TAgg '[] cols new) -> + TypedDataFrame cols -> + TypedDataFrame (Append cols (Reverse new)) +deriveMany build (TDF df) = + unsafeFreeze (D.deriveMany (taggToNamedExprs (build TAggNil)) df) + impute :: forall name a cols. ( KnownSymbol name @@ -336,6 +363,22 @@ symmetricDifference :: TypedDataFrame cols -> TypedDataFrame cols -> TypedDataFrame cols symmetricDifference (TDF a) (TDF b) = TDF (DS.symmetricDifference a b) +-- | The frame with no rows and no columns. +empty :: TypedDataFrame '[] +empty = TDF D.empty + +-- | True when the frame has no rows or no columns. +null :: TypedDataFrame cols -> Bool +null (TDF df) = D.null df + +-- | Thread the frame through a schema-preserving step for each element. +fold :: + (a -> TypedDataFrame cols -> TypedDataFrame cols) -> + [a] -> + TypedDataFrame cols -> + TypedDataFrame cols +fold f xs df = L.foldl' (flip f) df xs + dimensions :: TypedDataFrame cols -> (Int, Int) dimensions (TDF df) = D.dimensions df diff --git a/dataframe-parquet/src/DataFrame/IO/Parquet.hs b/dataframe-parquet/src/DataFrame/IO/Parquet.hs index 1a67a37a..4c677683 100644 --- a/dataframe-parquet/src/DataFrame/IO/Parquet.hs +++ b/dataframe-parquet/src/DataFrame/IO/Parquet.hs @@ -899,7 +899,7 @@ applySelectedColumns opts df = applyPredicate :: ParquetReadOptions -> DataFrame -> DataFrame applyPredicate opts df = - maybe df (`DS.filterWhere` df) (predicate opts) + maybe df (`DS.filter` df) (predicate opts) applySafeRead :: ParquetReadOptions -> DataFrame -> DataFrame applySafeRead opts df diff --git a/dataframe-parquet/src/DataFrame/Typed/IO/Parquet.hs b/dataframe-parquet/src/DataFrame/Typed/IO/Parquet.hs index b1a2e597..d98a4da6 100644 --- a/dataframe-parquet/src/DataFrame/Typed/IO/Parquet.hs +++ b/dataframe-parquet/src/DataFrame/Typed/IO/Parquet.hs @@ -8,15 +8,22 @@ module DataFrame.Typed.IO.Parquet ( readParquetWithError, readParquetWithOpts, readParquetFiles, + readParquetFilesWithOpts, + writeParquet, + writeParquetWithOptions, ) where import Control.Applicative ((<|>)) import Control.Exception (SomeException, try) import qualified Data.Text as T -import DataFrame.IO.Parquet (ParquetReadOptions (..), defaultParquetReadOptions) +import DataFrame.IO.Parquet ( + ParquetReadOptions (..), + ParquetWriteOptions, + defaultParquetReadOptions, + ) import qualified DataFrame.IO.Parquet as Parquet -import DataFrame.Typed.Freeze (freezeOrThrow, freezeWithError) +import DataFrame.Typed.Freeze (freezeOrThrow, freezeWithError, thaw) import DataFrame.Typed.Schema (KnownSchema, schemaColumnNames) import DataFrame.Typed.Types (TypedDataFrame) @@ -58,8 +65,23 @@ readParquetWithOpts opts path = -- | Read a directory\/glob of Parquet files into a typed DataFrame. readParquetFiles :: forall cols. (KnownSchema cols) => FilePath -> IO (TypedDataFrame cols) -readParquetFiles path = - Parquet.readParquetFilesWithOpts - (schemaOptions @cols defaultParquetReadOptions) - path +readParquetFiles = readParquetFilesWithOpts @cols defaultParquetReadOptions + +-- | Read a directory\/glob of Parquet files with custom options. +readParquetFilesWithOpts :: + forall cols. + (KnownSchema cols) => + ParquetReadOptions -> FilePath -> IO (TypedDataFrame cols) +readParquetFilesWithOpts opts path = + Parquet.readParquetFilesWithOpts (schemaOptions @cols opts) path >>= freezeOrThrow @cols + +-- | Write a typed DataFrame to a Parquet file. +writeParquet :: FilePath -> TypedDataFrame cols -> IO () +writeParquet path = Parquet.writeParquet path . thaw + +-- | Write a typed DataFrame to a Parquet file with custom options. +writeParquetWithOptions :: + ParquetWriteOptions -> FilePath -> TypedDataFrame cols -> IO () +writeParquetWithOptions opts path = + Parquet.writeParquetWithOptions opts path . thaw diff --git a/dataframe-persistent/README.md b/dataframe-persistent/README.md index c76b0673..26f1e380 100644 --- a/dataframe-persistent/README.md +++ b/dataframe-persistent/README.md @@ -175,7 +175,7 @@ A filter on a column that doesn't exist (or has the wrong type) is a compile err runtime surprise: ```haskell -D.toMarkdown' . DT.thaw . DT.filterWhere (DT.col @"Name" .==. DT.lit (Just "Accept")) <$> artists +D.toMarkdown' . DT.thaw . DT.filter (DT.col @"Name" .==. DT.lit (Just "Accept")) <$> artists ``` > diff --git a/dataframe-persistent/docs/base_scripts/base_readme.md b/dataframe-persistent/docs/base_scripts/base_readme.md index 1d7ead7e..3fed4f58 100644 --- a/dataframe-persistent/docs/base_scripts/base_readme.md +++ b/dataframe-persistent/docs/base_scripts/base_readme.md @@ -119,7 +119,7 @@ A filter on a column that doesn't exist (or has the wrong type) is a compile err runtime surprise: ```haskell -D.toMarkdown' . DT.thaw . DT.filterWhere (DT.col @"Name" .==. DT.lit (Just "Accept")) <$> artists +D.toMarkdown' . DT.thaw . DT.filter (DT.col @"Name" .==. DT.lit (Just "Accept")) <$> artists ``` Because the database is a value, reading the same table from two sources to join them is just two diff --git a/dataframe-persistent/tests/PersistentTests.hs b/dataframe-persistent/tests/PersistentTests.hs index aeec28c1..6ab4d789 100644 --- a/dataframe-persistent/tests/PersistentTests.hs +++ b/dataframe-persistent/tests/PersistentTests.hs @@ -193,7 +193,7 @@ testDataFrameOperations = TestCase $ withTestDb $ do df <- fromPersistent @TestUser [] liftIO $ do -- Filter operation - let youngUsers = DF.filter test_user_age (< 30) df + let youngUsers = DF.filterBy (< 30) test_user_age df assertEqual "Young users count" 2 (nRows youngUsers) -- Sort operation diff --git a/dataframe.cabal b/dataframe.cabal index bc44f6a1..30887288 100644 --- a/dataframe.cabal +++ b/dataframe.cabal @@ -75,7 +75,8 @@ library -- gets added to build-depends but never re-exported. ghc-options: -Werror=unused-packages exposed-modules: DataFrame, - DataFrame.Typed + DataFrame.Typed, + DataFrame.Typed.Interop reexported-modules: DataFrame.Functions, DataFrame.Synthesis, DataFrame.Display.Web.Plot, @@ -114,6 +115,7 @@ library DataFrame.Boosting.AdaBoost, DataFrame.Metrics, DataFrame.IO.JSON, + DataFrame.Typed.IO.JSON, DataFrame.Expr.Serialize, DataFrame.IR.ExprJson, DataFrame.Typed.Types, @@ -124,6 +126,7 @@ library DataFrame.Typed.Join, DataFrame.Typed.Aggregate, DataFrame.Typed.Expr, + DataFrame.Typed.Functions, DataFrame.Typed.Record, DataFrame.Typed.Generic build-depends: base >= 4 && <5, @@ -133,7 +136,8 @@ library dataframe-operations >= 2.6 && < 2.7, dataframe-parsing >= 2.2 && < 2.3, dataframe-viz >= 1.3.2 && < 1.4, - dataframe-learn >= 2.5 && < 2.6 + dataframe-learn >= 2.5 && < 2.6, + text >= 2.1 && < 3 if !flag(no-csv) reexported-modules: DataFrame.IO.CSV, @@ -341,6 +345,7 @@ test-suite tests Operations.Record, LazyParquet, LazyParity, + Parity.FourWay, LazyProjection, Parquet, ParquetTestData, diff --git a/docs/base/haskell_for_data_analysis.md b/docs/base/haskell_for_data_analysis.md index 9e2c0f91..4d61b636 100644 --- a/docs/base/haskell_for_data_analysis.md +++ b/docs/base/haskell_for_data_analysis.md @@ -370,14 +370,14 @@ pipelined = execFrameM housing $ do roomsPerHousehold <- deriveM "rooms_per_household" (total_rooms / households) meanBeds <- inspectM (D.meanMaybe total_bedrooms) totalBedrooms <- imputeM total_bedrooms meanBeds - filterWhereM (isExpensive .&&. roomsPerHousehold .>=. 7 .&&. totalBedrooms .>=. 200) + filterM (isExpensive .&&. roomsPerHousehold .>=. 7 .&&. totalBedrooms .>=. 200) TIO.putStrLn $ D.toMarkdown $ D.take 5 pipelined ``` `$(D.declareColumnsFromCsvFile path)` generates compile-time column bindings by reading the CSV header at splice time — no live dataframe needs to be in scope, unlike `$(D.declareColumns df)` which requires a bound frame. -Inside the `do`-block, `<-` binds the typed `Expr` returned by each step; those expressions can be reused in later steps (e.g. `isExpensive` in the final `filterWhereM`) without any extra plumbing. +Inside the `do`-block, `<-` binds the typed `Expr` returned by each step; those expressions can be reused in later steps (e.g. `isExpensive` in the final `filterM`) without any extra plumbing. **Extracting results alongside the final frame** — `runFrameM` returns a pair `(a, DataFrame)`: @@ -397,7 +397,7 @@ The three exit functions: | `runFrameM df block` | `(a, DataFrame)` | you need extracted values *and* the frame | | `evalFrameM df block` | `a` | you only need the extracted values | -Available monadic operations: `deriveM`, `imputeM`, `filterWhereM`, `inspectM`, `renameM`. +Available monadic operations: `deriveM`, `imputeM`, `filterM`, `inspectM`, `renameM`. --- @@ -500,7 +500,7 @@ TIO.putStrLn $ D.toMarkdown (D.take 5 withIncomeCast) ### Streaming large files with DataFrame.Lazy -All of the above reads the entire file into memory upfront. For datasets larger than available RAM, `DataFrame.Lazy` offers a pull-based streaming executor: operations build a logical plan tree and nothing is read from disk until `runDataFrame` is called. The optimizer pushes `filter` predicates down to the scan, so unneeded rows are discarded before any column is allocated. +All of the above reads the entire file into memory upfront. For datasets larger than available RAM, `DataFrame.Lazy` offers a pull-based streaming executor: operations build a logical plan tree and nothing is read from disk until `run` is called. The optimizer pushes `filter` predicates down to the scan, so unneeded rows are discarded before any column is allocated. The Lazy path requires an explicit schema — there is no inference. Build one with `schemaType @T`: @@ -533,7 +533,7 @@ lazyQuery = |> L.select ["longitude","latitude","median_house_value","ocean_proximity"] |> L.take 10 -lazyResult <- L.runDataFrame lazyQuery +lazyResult <- L.run lazyQuery TIO.putStrLn $ D.toMarkdown lazyResult ``` @@ -548,7 +548,7 @@ Key API: | `L.fromDataFrame df` | lift an eager frame into the lazy plan | | `L.filter`, `L.select`, `L.derive` | lazy counterparts of the eager operations | | `L.take`, `L.sortBy`, `L.groupBy`, `L.join` | same semantics, deferred execution | -| `L.runDataFrame plan` | execute the plan and materialise a `DataFrame` | +| `L.run plan` | execute the plan and materialise a `DataFrame` | The optimizer automatically reorders `filter` before `select` and prunes columns that are not referenced downstream — columns excluded by `L.select` are never allocated at the scan level. @@ -1083,6 +1083,7 @@ The trade-off is more explicit upfront ceremony (a TH splice, a freeze call, and ```haskell import qualified DataFrame.Typed as DT +import qualified DataFrame.Typed.Functions as DTF $(DT.deriveSchemaFromCsvFile "Housing" "../data/housing.csv") ``` @@ -1132,7 +1133,7 @@ TIO.putStrLn $ D.toMarkdown (DT.thaw typedResult) ```haskell typedGrouped = thousing |> DT.groupBy @'["ocean_proximity"] - |> DT.aggregate (DT.as @"count" (DT.count (DT.col @"median_house_value"))) + |> DT.aggregate (DT.as @"count" (DTF.count (DT.col @"median_house_value"))) TIO.putStrLn $ D.toMarkdown (DT.thaw typedGrouped) ``` diff --git a/docs/base_scripts/base_readme.md b/docs/base_scripts/base_readme.md index 68b94af0..87bb93d2 100644 --- a/docs/base_scripts/base_readme.md +++ b/docs/base_scripts/base_readme.md @@ -188,7 +188,7 @@ $(D.declareColumnsFromCsvFile "./data/housing.csv") df <- D.readCsv "./data/housing.csv" df |> D.derive "rooms_per_household" (total_rooms / households) - |> D.filterWhere (median_income .>. 5) + |> D.filter (median_income .>. 5) |> D.groupBy ["ocean_proximity"] |> D.aggregate [F.mean median_house_value `as` "avg_value"] |> D.toMarkdown' @@ -200,7 +200,7 @@ Compare this to the manual version which requires spelling out every column name -- 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) + |> D.filter (F.col @Double "median_income" .>. F.lit 5) |> D.take 5 |> D.toMarkdown' ``` @@ -261,7 +261,7 @@ D.toRecords ordersDf :: Either Text [Order] And the typed bridge — `[Order]` to `TypedDataFrame OrderSchema` and back: ```haskell -DT.thaw (DT.fromRecordsTyped xs :: DT.TypedDataFrame OrderSchema) |> D.toMarkdown' +DT.thaw (DT.fromRecords xs :: DT.TypedDataFrame OrderSchema) |> D.toMarkdown' ``` Field names are translated `camelCase → snake_case` by default; override @@ -287,7 +287,7 @@ $(D.deriveSchemaValues ''Order) orders :: IO D.DataFrame orders = do raw <- D.readCsvWithSchema orderSchema "./data/orders.csv" - pure (D.filter orderAmount (> 100) raw) + pure (D.filterBy (> 100) orderAmount raw) ``` Each record field gets a typed accessor named ``, @@ -332,7 +332,7 @@ case DT.freeze @EmployeeSchema employees of Nothing -> "Schema mismatch!" Just tdf -> tdf |> DT.derive @"bonus" (DT.col @"salary" * DT.lit 0.1) - |> DT.filterWhere (DT.col @"salary" DT..>. DT.lit 50000) + |> DT.filter (DT.col @"salary" DT..>. DT.lit 50000) |> DT.select @'["name", "bonus"] |> DT.thaw |> D.toMarkdown' @@ -342,11 +342,11 @@ case DT.freeze @EmployeeSchema employees of ```text -- Typo in column name -> compile error -tdf |> DT.filterWhere (DT.col @"slary" DT..>. DT.lit 50000) +tdf |> DT.filter (DT.col @"slary" DT..>. DT.lit 50000) -- error: Column 'slary' not found in schema -- Wrong type -> compile error -tdf |> DT.filterWhere (DT.col @"name" DT..>. DT.lit 50000) +tdf |> DT.filter (DT.col @"name" DT..>. DT.lit 50000) -- error: Couldn't match type 'Text' with 'Double' ``` @@ -369,7 +369,7 @@ DT.thaw (DT.filterAllJust stdf |> DT.derive @"scaled" (DT.col @"score" * DT.lit ## 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: +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 operations as the eager API, and `run` passes it through an optimizer before streaming results batch-by-batch: ```haskell import qualified DataFrame.Lazy as L @@ -388,7 +388,7 @@ housingSchema = makeSchema , ("ocean_proximity", schemaType @Text) ] -lazyResult <- L.runDataFrame $ +lazyResult <- L.run $ L.scanCsv housingSchema "./data/housing.csv" |> L.filter (F.col @Double "median_income" .>. F.lit 5) |> L.derive "value_per_income" diff --git a/docs/coming_from_other_implementations.md b/docs/coming_from_other_implementations.md index b6759a9b..b7a9f989 100644 --- a/docs/coming_from_other_implementations.md +++ b/docs/coming_from_other_implementations.md @@ -24,7 +24,7 @@ same underlying `DataFrame` type at runtime. | **Typed** | x | x | Production code, libraries, anywhere a schema change should be a compile error | **Untyped** is what most of this guide uses: expressions built with `F.col @Type "name"` and -operations like `D.derive`, `D.filterWhere`, `D.groupBy`. Column names are strings — typos only +operations like `D.derive`, `D.filter`, `D.groupBy`. Column names are strings — typos only surface at runtime. **Frame monad** (`FrameM`) is a thin state monad over `DataFrame`. `deriveM` returns the @@ -359,8 +359,8 @@ Since we don't have row indexes, we filter by actual values using predicates: ```haskell dataframe> :{ -ghci| df' |> D.filter "date" (\d -> d >= (fromGregorian 2013 01 02) -ghci| && d <= (fromGregorian 2013 01 04)) +ghci| df' |> D.filterBy (\d -> d >= (fromGregorian 2013 01 02) +ghci| && d <= (fromGregorian 2013 01 04)) (F.col @Day "date") ghci| |> D.select ["A", "B"] ghci| :} -------------------- @@ -394,12 +394,12 @@ import qualified DataFrame.Functions as F import DataFrame.Expression.Operators df |> D.derive "doubled_A" (F.col @Double "A" * F.lit 2) - |> D.filterWhere (F.col @Double "doubled_A" .>. F.lit 1.0) + |> D.filter (F.col @Double "doubled_A" .>. F.lit 1.0) |> D.select ["date", "doubled_A"] ``` Straightforward, but the string `"doubled_A"` appears twice. If you rename the column in `derive` -but forget to update `filterWhere`, the error only appears when the code runs. +but forget to update `filter`, the error only appears when the code runs. #### Frame monad @@ -408,16 +408,16 @@ import DataFrame.Monad execFrameM df $ do doubledA <- deriveM "doubled_A" (F.col @Double "A" * F.lit 2) - filterWhereM (doubledA .>. F.lit 1.0) + filterM (doubledA .>. F.lit 1.0) modifyM (D.select ["date", "doubled_A"]) ``` `deriveM` returns `doubledA :: Expr Double` — the live expression for the column that was just -added. The subsequent `filterWhereM` uses that value directly rather than looking the column +added. The subsequent `filterM` uses that value directly rather than looking the column up by name. **Why it's better**: the string `"doubled_A"` appears exactly once. A rename refactoring only -touches the `deriveM` call; the `filterWhereM` argument follows automatically because it holds a +touches the `deriveM` call; the `filterM` argument follows automatically because it holds a reference to the expression, not a copy of the name. #### Typed API @@ -425,6 +425,7 @@ reference to the expression, not a copy of the name. ```haskell {-# LANGUAGE DataKinds, TypeApplications #-} import qualified DataFrame.Typed as T +import qualified DataFrame.Typed.Functions as TF import Data.Time.Calendar (Day) type MySchema = '[ '("date", Day) @@ -438,12 +439,12 @@ case T.freeze @MySchema df of Nothing -> error "schema mismatch at startup" Just tdf -> tdf |> T.derive @"doubled_A" (T.col @"A" * T.lit 2) - |> T.filterWhere (T.col @"doubled_A" .>. T.lit 1.0) + |> T.filter (T.col @"doubled_A" .>. T.lit 1.0) |> T.select @'["date", "doubled_A"] ``` **Why it's better**: `T.col @"A"` is a compile-time error if the column `"A"` does not exist in -`MySchema` or has the wrong type. The same applies to `T.col @"doubled_A"` in `filterWhere` — if +`MySchema` or has the wrong type. The same applies to `T.col @"doubled_A"` in `filter` — if you accidentally filter before deriving, or use the wrong type annotation, the code does not compile. Typos in column names, wrong aggregation types, and schema-breaking refactors are all caught before the program runs. @@ -756,7 +757,7 @@ print(result) let bornBefore1990 d = let (y, _, _) = toGregorian d in y < 1990 -df_csv |> D.filter "birthdate" bornBefore1990 +df_csv |> D.filterBy bornBefore1990 (F.col @Day "birthdate") ``` **Output:** @@ -772,7 +773,7 @@ df_csv |> D.filter "birthdate" bornBefore1990 Daniel Donovan | 1981-04-30 | 83.1 | 1.75 ``` -The predicate function receives individual values from the named column and returns `True` to keep the row. +The predicate function receives individual values from the column and returns `True` to keep the row. #### Multiple Filter Conditions @@ -792,8 +793,8 @@ let year d = let (y, _, _) = toGregorian d in y between lo hi val = val >= lo && val <= hi df_csv - |> D.filter "birthdate" (between 1982 1996 . year) - |> D.filter "height" (> 1.7) + |> D.filterBy (between 1982 1996 . year) (F.col @Day "birthdate") + |> D.filter (F.col @Double "height" .>. 1.7) ``` **Output:** @@ -809,11 +810,11 @@ df_csv Each filter operates on the DataFrame sequentially, and the types ensure you can't filter on non-existent columns. -Alternatively, you can use `filterWhere` with boolean expression combinations: +Alternatively, you can use `filter` with boolean expression combinations: ```haskell df_csv - |> D.filterWhere ((F.col @Int "birth_year" .>=. 1982) + |> D.filter ((F.col @Int "birth_year" .>=. 1982) .&&. (F.col @Int "birth_year" .<=. 1996) .&&. (F.col @Double "height" .>. 1.7)) ``` @@ -1005,6 +1006,7 @@ Here the schema is declared upfront and every column access is validated at comp ```haskell {-# LANGUAGE DataKinds, TypeApplications #-} import qualified DataFrame.Typed as T +import qualified DataFrame.Typed.Functions as TF import Data.Time.Calendar (Day) type BirthdateSchema = '[ '("name", T.Text) @@ -1019,14 +1021,14 @@ type BirthdateSchema = '[ '("name", T.Text) example :: T.TypedDataFrame BirthdateSchema -> IO () example tdf = do let result = T.aggregate - ( T.as @"avg_weight" (T.mean (T.col @"weight")) - . T.as @"avg_height" (T.mean (T.col @"height")) - . T.as @"names" (T.collect (T.col @"name")) + ( T.as @"avg_weight" (TF.mean (T.col @"weight")) + . T.as @"avg_height" (TF.mean (T.col @"height")) + . T.as @"names" (TF.collect (T.col @"name")) ) (T.groupBy @'["decade"] tdf') tdf' = tdf - |> T.derive @"name" (T.lift firstName (T.col @"name")) - |> T.derive @"decade" (T.lift decade (T.col @"birthdate")) + |> T.derive @"name" (TF.lift firstName (T.col @"name")) + |> T.derive @"decade" (TF.lift decade (T.col @"birthdate")) |> T.exclude @'["birthdate"] print (T.thaw result) where @@ -1034,7 +1036,7 @@ example tdf = do firstName = head . T.split (== ' ') ``` -**Why it's better**: `T.as @"avg_weight" (T.mean (T.col @"weight"))` is checked in two ways at +**Why it's better**: `T.as @"avg_weight" (TF.mean (T.col @"weight"))` is checked in two ways at compile time — `"weight"` must exist in the schema with type `Double`, and the result column `"avg_weight"` will have type `Double` in the output schema. A wrong output-type annotation causes a type error before the program runs. @@ -1077,7 +1079,7 @@ mySchema = [ ("name", schemaType @T.Text) ] result :: IO DataFrame -result = L.runDataFrame $ +result = L.run $ L.scanCsv mySchema "data.csv" |> L.filter (F.col @Double "height" .>. F.lit 1.7) |> L.select ["name", "weight", "height"] @@ -1093,7 +1095,7 @@ workflows. Its type inference is automatic — you rarely need to declare schem - **Schema required upfront**: `mySchema` tells the optimizer the column names and types before touching the file. The optimizer knows which columns to read and what types to allocate in each batch — no runtime inference pass needed. -- **Optimizer runs before IO**: `L.runDataFrame` calls the rule-based optimizer first, then +- **Optimizer runs before IO**: `L.run` calls the rule-based optimizer first, then streams. Dead columns (those not referenced downstream) are eliminated at the scan level — bytes never leave disk. Filter predicates are pushed into `ScanConfig` so each batch is filtered on arrival. @@ -1108,7 +1110,7 @@ workflows. Its type inference is automatic — you rarely need to declare schem ```haskell result :: IO DataFrame -result = L.runDataFrame $ +result = L.run $ L.scanParquet mySchema "warehouse/events.parquet" |> L.filter (F.col @T.Text "country" .==. F.lit "US") |> L.select ["event_id", "country", "revenue"] @@ -1154,7 +1156,7 @@ starwars %>% import qualified Data.Text as T starwars - |> D.filterWhere (F.col @Text "species" .==. "Droid") + |> D.filter (F.col @Text "species" .==. "Droid") ``` **Output (truncated for readability):** @@ -1346,7 +1348,7 @@ starwars |> D.aggregate [ F.mean mass `as` "mean_mass" , F.count mass `as` "count" ] - |> D.filterWhere ((F.col @Int "count" .>. 1) .&&. (F.col @Double "mean_mass" .>. 50)) + |> D.filter ((F.col @Int "count" .>. 1) .&&. (F.col @Double "mean_mass" .>. 50)) ``` **Output:** @@ -1370,7 +1372,7 @@ starwars **Key points:** 1. We use `aggregate` with a list of aggregation expressions -2. `filterWhere` allows filtering based on multiple column conditions using `.&&.` (AND) and `.||.` (OR) +2. `filter` allows filtering based on multiple column conditions using `.&&.` (AND) and `.||.` (OR) 3. Type annotations ensure we're comparing the right types ### BMI + Group Pipeline — Frame Monad Version @@ -1385,12 +1387,12 @@ execFrameM starwars $ do bmiCol <- deriveM "bmi" (F.col @Double "mass" / F.pow (F.col @Double "height" / F.lit 100) 2) - filterWhereM (bmiCol .>. F.lit 20.0) + filterM (bmiCol .>. F.lit 20.0) modifyM (D.select ["name", "height", "mass", "bmi"]) ``` **Why it's better**: `D.filterAllJust` drops rows with any `Nothing` in one call — no separate -`filterJust` per column. `deriveM` returns the column expression directly; `filterWhereM` uses +`filterJust` per column. `deriveM` returns the column expression directly; `filterM` uses it without repeating the name `"bmi"`. The pipeline is a sequence of plain `do` steps that reads like prose. @@ -1402,6 +1404,7 @@ in the result, so any downstream code that still treats them as optional fails a ```haskell {-# LANGUAGE DataKinds, TypeApplications #-} import qualified DataFrame.Typed as T +import qualified DataFrame.Typed.Functions as TF type StarwarsSchema = '[ '("name", T.Text) @@ -1420,7 +1423,7 @@ example tdf = (T.col @"mass" / (T.col @"height" / T.lit 100) ^ 2) stripped in withBmi - |> T.filterWhere (T.col @"bmi" .>. T.lit 20.0) + |> T.filter (T.col @"bmi" .>. T.lit 20.0) |> T.select @'["name", "height", "mass", "bmi"] ``` @@ -1666,7 +1669,7 @@ purchases |> ```haskell let med = D.median (F.col @Double "amount") df -df |> D.filterWhere (F.col @Double "amount" .<=. F.lit (med * 10)) +df |> D.filter (F.col @Double "amount" .<=. F.lit (med * 10)) |> D.derive "net" (F.col @Int "amount" - F.col @Int "discount") |> D.groupBy ["country"] |> D.aggregate [F.sum (F.col @Int "net") `as` "total"] @@ -1718,7 +1721,7 @@ let medians = df |> D.aggregate [F.median (F.col @Double "amount") `as` "country_median"] D.innerJoin ["country"] df medians - |> D.filterWhere (F.col @Double "amount" .<=. F.col @Double "country_median" * 10) + |> D.filter (F.col @Double "amount" .<=. F.col @Double "country_median" * 10) |> D.derive "net" (F.col @Int "amount" - F.col @Int "discount") |> D.groupBy ["country"] |> D.aggregate [F.sum (F.col @Int "net") `as` "total"] @@ -1737,6 +1740,7 @@ the vinyl overhead: import qualified DataFrame as D import qualified DataFrame.Typed as T +import qualified DataFrame.Typed.Functions as TF import Data.Text (Text) import DataFrame.Expression.Operators @@ -1753,7 +1757,7 @@ main = do Nothing -> putStrLn "Schema mismatch!" Just tdf -> do let result = T.aggregate - (T.as @"total" (T.sum (T.col @"amount"))) + (T.as @"total" (TF.sum (T.col @"amount"))) (T.groupBy @'["country"] tdf) print (T.thaw result) ``` @@ -1817,17 +1821,17 @@ tdf |> T.derive @"bonus" (T.col @"salary" * T.lit 0.1) **Untyped** ```haskell -df |> D.filterWhere (F.col @Double "salary" .>. F.lit 50000) +df |> D.filter (F.col @Double "salary" .>. F.lit 50000) ``` **Frame monad** ```haskell -execFrameM df $ filterWhereM (F.col @Double "salary" .>. F.lit 50000) +execFrameM df $ filterM (F.col @Double "salary" .>. F.lit 50000) ``` **Typed** ```haskell -tdf |> T.filterWhere (T.col @"salary" .>. T.lit 50000) +tdf |> T.filter (T.col @"salary" .>. T.lit 50000) ``` #### Grouping and aggregating @@ -1853,8 +1857,8 @@ execFrameM df $ do **Typed** ```haskell T.aggregate - ( T.as @"avg_salary" (T.mean (T.col @"salary")) - . T.as @"n" (T.count (T.col @"salary")) + ( T.as @"avg_salary" (TF.mean (T.col @"salary")) + . T.as @"n" (TF.count (T.col @"salary")) ) (T.groupBy @'["dept"] tdf) ``` @@ -1892,7 +1896,7 @@ let filled = T.impute @"col" (0 :: Double) tdf ```haskell import qualified DataFrame.Lazy as L -result <- L.runDataFrame $ +result <- L.run $ L.scanCsv mySchema "large_file.csv" |> L.filter (F.col @Double "revenue" .>. F.lit 1000) |> L.select ["id", "region", "revenue"] diff --git a/docs/cookbook.md b/docs/cookbook.md index 5a8d3357..d90cdb78 100644 --- a/docs/cookbook.md +++ b/docs/cookbook.md @@ -21,14 +21,14 @@ The `:declareColumns` command makes column names available as variables in your ## Filtering Data -One of the most fundamental operations in data analysis is filtering - selecting rows that meet certain criteria. In Haskell's DataFrame library, we use the `filterWhere` function combined with comparison operators. +One of the most fundamental operations in data analysis is filtering - selecting rows that meet certain criteria. In Haskell's DataFrame library, we use the `filter` function combined with comparison operators. ### Basic Comparisons -The `filterWhere` function takes a boolean expression and returns only the rows where that expression evaluates to true. For example, to find rows where a column equals a specific value, we use the `.==.` operator (same-type, non-nullable) or `.==` (nullable-aware): +The `filter` function takes a boolean expression and returns only the rows where that expression evaluates to true. For example, to find rows where a column equals a specific value, we use the `.==.` operator (same-type, non-nullable) or `.==` (nullable-aware): ```haskell -df |> D.filterWhere (columnName .==. value) +df |> D.filter (columnName .==. value) ``` The pipe operator `|>` allows us to chain operations in a readable left-to-right style, similar to Unix pipes. @@ -42,7 +42,7 @@ Query all columns for a city with the ID 1661. ### Solution ```haskell -dataframe> df |> D.filterWhere (id .==. 1661) +dataframe> df |> D.filter (id .==. 1661) ----------------------------------------------------- id | name | country_code | district | population ------|--------|--------------|----------|----------- @@ -60,7 +60,7 @@ Query all columns of every Japanese city. The `country_code` for Japan is "JPN". ### Solution ```haskell -dataframe> df |> D.filterWhere (country_code .==. "JPN") +dataframe> df |> D.filter (country_code .==. "JPN") -------------------------------------------------------- id | name | country_code | district | population ------|----------|--------------|-----------|----------- @@ -84,7 +84,7 @@ Often you'll need to filter on multiple conditions simultaneously. You can combi For example, to find cities with large populations in a specific country: ```haskell -df |> D.filterWhere ((population .>. 100000) .&&. (country_code .==. "USA")) +df |> D.filter ((population .>. 100000) .&&. (country_code .==. "USA")) ``` **Exercise 3: Basic filtering (cont)** @@ -101,7 +101,7 @@ Query all columns for all American cities in city dataframe with: ```haskell dataframe> D.readCsv "./data/country.csv" dataframe> :declareColumns df -dataframe> df |> D.filterWhere ((population .>. 100000) .&&. (country_code .==. "USA")) +dataframe> df |> D.filter ((population .>. 100000) .&&. (country_code .==. "USA")) -------------------------------------------------------------- id | name | country_code | district | population ------|---------------|--------------|------------|----------- @@ -183,7 +183,7 @@ dataframe> df |> D.select [F.name name] |> D.take 5 The real power of these operations comes from chaining them together. You can filter rows and then select specific columns (or vice versa) to get exactly the data you need: ```haskell -df |> D.filterWhere (condition) |> D.select [columns] |> D.take n +df |> D.filter (condition) |> D.select [columns] |> D.take n ``` The order of operations matters - filtering first reduces the data before selection, which can be more efficient. @@ -197,7 +197,7 @@ Query the names of all the Japanese cities and show only the first 5 results. ### Solution ```haskell -dataframe> df |> D.filterWhere (country_code .==. "JPN") |> D.select [F.name name] |> D.take 5 +dataframe> df |> D.filter (country_code .==. "JPN") |> D.select [F.name name] |> D.take 5 --------- name --------- @@ -251,7 +251,7 @@ Query a list of city names for cities that have an even ID number. Show the resu ### Solution ```haskell -dataframe> df |> D.filterWhere (F.lift even id) |> D.select [F.name city] |> D.distinct +dataframe> df |> D.filter (F.lift even id) |> D.select [F.name city] |> D.distinct ---------------------- city ---------------------- @@ -326,7 +326,7 @@ dataframe> (letterSort D.Desc) <> (letterSort D.Asc) One of the strengths of working with dataframes in Haskell is the ability to use any Haskell function in your queries. The `F.lift` function allows you to apply regular Haskell functions to DataFrame columns. This means you can use string functions, mathematical operations, or even your own custom logic: ```haskell -df |> D.filterWhere (F.lift customFunction columnName) +df |> D.filter (F.lift customFunction columnName) ``` This enables sophisticated filtering that goes beyond simple comparisons. For example, you can check string prefixes, perform calculations, or apply complex business logic. @@ -340,7 +340,7 @@ Query the list of city names starting with vowels (i.e., a, e, i, o, or u). Your ### Solution ```haskell -dataframe> df |> D.select [F.name city] |> D.filterWhere (F.lift (\c -> any (`T.isPrefixOf` (T.toLower c)) ["a", "e", "i", "o", "u"]) city) |> D.take 5 +dataframe> df |> D.select [F.name city] |> D.filter (F.lift (\c -> any (`T.isPrefixOf` (T.toLower c)) ["a", "e", "i", "o", "u"]) city) |> D.take 5 ---------- city ---------- @@ -500,6 +500,7 @@ module Main where import qualified DataFrame as D import qualified DataFrame.Functions as F import qualified DataFrame.Typed as DT +import qualified DataFrame.Typed.Functions as DTF import DataFrame.Expression.Operators ((|>)) import DataFrame.Typed ((.>.)) @@ -537,16 +538,16 @@ main = do submissions' |> DT.groupBy @'["challenge_id"] |> DT.aggregate - ( (DT.sum (DT.col @"total_submissions") |> DT.as @"total_submissions") - . (DT.sum (DT.col @"total_accepted_submissions") + ( (DTF.sum (DT.col @"total_submissions") |> DT.as @"total_submissions") + . (DTF.sum (DT.col @"total_accepted_submissions") |> DT.as @"total_accepted_submissions") ) let viewTotals = views' |> DT.groupBy @'["challenge_id"] |> DT.aggregate - ( (DT.sum (DT.col @"total_views") |> DT.as @"total_views") - . (DT.sum (DT.col @"total_unique_views") |> DT.as @"total_unique_views") + ( (DTF.sum (DT.col @"total_views") |> DT.as @"total_views") + . (DTF.sum (DT.col @"total_unique_views") |> DT.as @"total_unique_views") ) print $ contestsWithColleges @@ -565,7 +566,7 @@ main = do |> DT.impute @"total_views" (0 :: Int) |> DT.impute @"total_submissions" (0 :: Int) |> DT.impute @"total_accepted_submissions" (0 :: Int) - |> DT.filterWhere + |> DT.filter ( DT.col @"total_unique_views" + DT.col @"total_views" + DT.col @"total_submissions" @@ -577,7 +578,7 @@ main = do You've now learned the fundamental operations for working with dataframes in Haskell: -- **Filtering** with `filterWhere` to select rows based on conditions +- **Filtering** with `filter` to select rows based on conditions - **Selecting** with `select` to choose specific columns - **Limiting** with `take` to control output size - **Removing duplicates** with `distinct` diff --git a/docs/haskell_for_data_analysis.md b/docs/haskell_for_data_analysis.md index 12c95e1b..6ade2130 100644 --- a/docs/haskell_for_data_analysis.md +++ b/docs/haskell_for_data_analysis.md @@ -660,7 +660,7 @@ pipelined = execFrameM housing $ do roomsPerHousehold <- deriveM "rooms_per_household" (total_rooms / households) meanBeds <- inspectM (D.meanMaybe total_bedrooms) totalBedrooms <- imputeM total_bedrooms meanBeds - filterWhereM (isExpensive .&&. roomsPerHousehold .>=. 7 .&&. totalBedrooms .>=. 200) + filterM (isExpensive .&&. roomsPerHousehold .>=. 7 .&&. totalBedrooms .>=. 200) TIO.putStrLn $ D.toMarkdown $ D.take 5 pipelined ``` @@ -677,7 +677,7 @@ TIO.putStrLn $ D.toMarkdown $ D.take 5 pipelined `$(D.declareColumnsFromCsvFile path)` generates compile-time column bindings by reading the CSV header at splice time (you don't need a live dataframe in scope) unlike `$(D.declareColumns df)` which requires a bound frame. -Inside the `do`-block, `<-` binds the typed `Expr` returned by each step; those expressions can be reused in later steps (e.g. `isExpensive` in the final `filterWhereM`) without any extra plumbing. +Inside the `do`-block, `<-` binds the typed `Expr` returned by each step; those expressions can be reused in later steps (e.g. `isExpensive` in the final `filterM`) without any extra plumbing. **Extracting results alongside the final frame** — `runFrameM` returns a pair `(a, DataFrame)`: @@ -701,7 +701,7 @@ The three exit functions: | `runFrameM df block` | `(a, DataFrame)` | you need extracted values *and* the frame | | `evalFrameM df block` | `a` | you only need the extracted values | -Available monadic operations: `deriveM`, `imputeM`, `filterWhereM`, `inspectM`, `renameM`. +Available monadic operations: `deriveM`, `imputeM`, `filterM`, `inspectM`, `renameM`. --- @@ -859,7 +859,7 @@ TIO.putStrLn $ D.toMarkdown (D.take 5 withIncomeCast) ### Streaming large files with DataFrame.Lazy -All of the above reads the entire file into memory upfront. For datasets larger than available RAM, `DataFrame.Lazy` offers a pull-based streaming executor: operations build a logical plan tree and nothing is read from disk until `runDataFrame` is called. The optimizer pushes `filter` predicates down to the scan, so unneeded rows are discarded before any column is allocated. +All of the above reads the entire file into memory upfront. For datasets larger than available RAM, `DataFrame.Lazy` offers a pull-based streaming executor: operations build a logical plan tree and nothing is read from disk until `run` is called. The optimizer pushes `filter` predicates down to the scan, so unneeded rows are discarded before any column is allocated. The Lazy path requires an explicit schema — there is no inference. Build one with `schemaType @T`: @@ -895,7 +895,7 @@ lazyQuery = |> L.select ["longitude","latitude","median_house_value","ocean_proximity"] |> L.take 10 -lazyResult <- L.runDataFrame lazyQuery +lazyResult <- L.run lazyQuery TIO.putStrLn $ D.toMarkdown lazyResult ``` @@ -911,7 +911,7 @@ Key API: | `L.fromDataFrame df` | lift an eager frame into the lazy plan | | `L.filter`, `L.select`, `L.derive` | lazy counterparts of the eager operations | | `L.take`, `L.sortBy`, `L.groupBy`, `L.join` | same semantics, deferred execution | -| `L.runDataFrame plan` | execute the plan and materialise a `DataFrame` | +| `L.run plan` | execute the plan and materialise a `DataFrame` | The optimizer automatically reorders `filter` before `select` and prunes columns that are not referenced downstream — columns excluded by `L.select` are never allocated at the scan level. @@ -1529,6 +1529,7 @@ The trade-off is more explicit upfront ceremony (a TH splice, a freeze call, and ```haskell import qualified DataFrame.Typed as DT +import qualified DataFrame.Typed.Functions as DTF $(DT.deriveSchemaFromCsvFile "Housing" "../data/housing.csv") ``` @@ -1586,7 +1587,7 @@ TIO.putStrLn $ D.toMarkdown (DT.thaw typedResult) ```haskell typedGrouped = thousing |> DT.groupBy @'["ocean_proximity"] - |> DT.aggregate (DT.as @"count" (DT.count (DT.col @"median_house_value"))) + |> DT.aggregate (DT.as @"count" (DTF.count (DT.col @"median_house_value"))) TIO.putStrLn $ D.toMarkdown (DT.thaw typedGrouped) ``` diff --git a/docs/persistent_integration.md b/docs/persistent_integration.md index 3973396a..7bb387f2 100644 --- a/docs/persistent_integration.md +++ b/docs/persistent_integration.md @@ -104,7 +104,7 @@ analyzeUsers = runSqlite "test.db" $ do liftIO $ do -- Filter - let youngUsers = DF.filter @Int "age" (< 30) df + let youngUsers = DF.filterBy (< 30) (F.col @Int "age") df -- Sort let sorted = DF.sortBy [DF.Asc "age"] df diff --git a/docs/quick_start.md b/docs/quick_start.md index ee9d5b6f..d1283748 100644 --- a/docs/quick_start.md +++ b/docs/quick_start.md @@ -78,7 +78,7 @@ main = do -- 2) creates a type safe reference to the newly created column. is_expensive <- deriveM "is_expensive" (median_house_value .>=. 500000) luxury <- deriveM "luxury" (is_expensive .&&. median_income .>=. 8) - filterWhereM luxury + filterM luxury ``` ### As a commandline tool diff --git a/docs/tutorial.md b/docs/tutorial.md index 1ef5dbff..4a61097c 100644 --- a/docs/tutorial.md +++ b/docs/tutorial.md @@ -80,7 +80,7 @@ This is a common idiom in languages like R and Julia, where it's sometimes calle ## 3) Filter rows with typed predicates -Let's say we want to filter all flowers whose petal width is greater than 6. We can use the `filterWhere` function. `filterWhere` takes a boolean expression (`Expr Bool`). Dataframes have an expression syntax (similar to Polars) that allow us to write logic on top of columns. +Let's say we want to filter all flowers whose petal width is greater than 6. We can use the `filter` function. `filter` takes a boolean expression (`Expr Bool`). Dataframes have an expression syntax (similar to Polars) that allow us to write logic on top of columns. There are a couple of core functions we can use to define columns: @@ -116,7 +116,7 @@ Armed with this knowledge, we can go back and filter all flowers with petal leng We can see from the sample we printed before that `petal.length` is of type `Double`. We write the expression as follows: ```haskell -df |> D.filterWhere (F.col @Double "petal.length" .>. F.lit @Double 6) +df |> D.filter (F.col @Double "petal.length" .>. F.lit @Double 6) |> D.select ["petal.width", "petal.length", "variety"] ``` @@ -126,7 +126,7 @@ Suppose we write out the wrong types in the expression. That is suppose we say t ```haskell -- This fails at runtime because "petal.length" is Double, not Int: -df |> D.filterWhere (F.col @Int "petal.length" .>. F.lit @Int 6) +df |> D.filter (F.col @Int "petal.length" .>. F.lit @Int 6) ``` ![Screenshot of filtering with type error](./_static/filter_wrong_type.png) @@ -154,7 +154,7 @@ and is re-exported from `DataFrame` for convenience.) In fact, we can make this even shorter. Since Haskell knows how to create any num instance from literals, we don't have to write `F.lit @Int 6`. It knows, from the context, to wrap `6` so it's an `Expr Double`. Thus, we can write: ```haskell -df |> D.filterWhere (petal_length .>. 6) +df |> D.filter (petal_length .>. 6) |> D.select ["petal.width", "petal.length", "variety"] ``` diff --git a/docs/using_dataframe_in_a_standalone_script.md b/docs/using_dataframe_in_a_standalone_script.md index 52cf846b..ccf2762b 100644 --- a/docs/using_dataframe_in_a_standalone_script.md +++ b/docs/using_dataframe_in_a_standalone_script.md @@ -147,7 +147,7 @@ main = do meanBedrooms = D.meanMaybe total_bedrooms dfWithRoomsPerHousehold dfWithTotalBedroomsImputed = D.impute total_bedrooms meanBedrooms dfWithRoomsPerHousehold - print $ dfWithTotalBedroomsImputed |> D.filterWhere (isExpensive .&&. roomsPerHousehold .>=. 7 .&&. (F.col @Double "total_bedrooms") .>=. 200) + print $ dfWithTotalBedroomsImputed |> D.filter (isExpensive .&&. roomsPerHousehold .>=. 7 .&&. (F.col @Double "total_bedrooms") .>=. 200) ``` Our code is now inundated with different dataframe variables and a bunch of pipeline management. Should we just revert back to using the old `derive` and keep track of names and types ourselves? Not at all. We can use a structure called a `FrameM` to remove this boilerplate while still keeping the guarantees we unlocked before. @@ -175,7 +175,7 @@ main = do roomsPerHousehold <- deriveM "rooms_per_household" (total_rooms / households) meanBedrooms <- inspectM (D.meanMaybe total_bedrooms) totalBedrooms <- imputeM total_bedrooms meanBedrooms - filterWhereM (isExpensive .&&. roomsPerHousehold .>=. 7 .&&.totalBedrooms .>=. 200) + filterM (isExpensive .&&. roomsPerHousehold .>=. 7 .&&.totalBedrooms .>=. 200) print df' ``` @@ -183,7 +183,7 @@ In the earlier version (without FrameM), every operation has to take an input da FrameM is a way to say: "For the next few lines, assume we're working on one evolving dataframe, and keep passing it along for me." You can think of it as a do-block where the dataframe is the implicit context. Inside the block, each step sees "the current dataframe" and can update it, but you don’t have to keep writing the plumbing yourself. -That’s what the M suffix is hinting at. When you see functions like deriveM, imputeM, or filterWhereM, the M is just a conventional signal that this version runs in a monad. And "monad", here, can be understood very simply: it’s a pattern for sequencing operations where each line may carry along some hidden context. In IO (like in the main function for example), the hidden context is "the outside world" (the console screen, system clocks, environment variables etc). In FrameM, the hidden context is "the current dataframe". +That’s what the M suffix is hinting at. When you see functions like deriveM, imputeM, or filterM, the M is just a conventional signal that this version runs in a monad. And "monad", here, can be understood very simply: it’s a pattern for sequencing operations where each line may carry along some hidden context. In IO (like in the main function for example), the hidden context is "the outside world" (the console screen, system clocks, environment variables etc). In FrameM, the hidden context is "the current dataframe". A nice way to read the FrameM part without doing "monads as a theory lesson" is to treat it like a mini language for dataframe steps. You write your steps in a natural order, and the monad takes care of passing the context to the next step. @@ -210,7 +210,7 @@ main = do print $ df' - |> D.filterWhere (isExpensive .&&. totalBedrooms .>=. 200) + |> D.filter (isExpensive .&&. totalBedrooms .>=. 200) ``` Now we can reuse the derived Exprs elsewhere - for example, in plotting, or further filtering with the regular dataframe API. @@ -253,6 +253,7 @@ To start, we need to define our exact data structure. In the full code example, import qualified DataFrame as D import qualified DataFrame.Typed as DT +import qualified DataFrame.Typed.Functions as DTF $(DT.deriveSchemaFromCsvFile "Housing" "./data/housing.csv") @@ -271,6 +272,7 @@ You can also define the schema by hand as follows: import qualified DataFrame as D import qualified DataFrame.Typed as DT +import qualified DataFrame.Typed.Functions as DTF type Housing = [ '("latitude", Double) , '("longitude", Double) @@ -296,6 +298,7 @@ When we read data from a file at runtime using D.readCsv, the initial dataframe import qualified DataFrame as D import qualified DataFrame.Typed as DT +import qualified DataFrame.Typed.Functions as DTF $(DT.deriveSchemaFromCsvFile "Housing" "./data/housing.csv") @@ -318,6 +321,7 @@ Once your dataframe is safely frozen, you can perform transformations with compl import qualified DataFrame as D import qualified DataFrame.Typed as DT +import qualified DataFrame.Typed.Functions as DTF import DataFrame.Expression.Operators ((|>)) @@ -349,6 +353,7 @@ This strict tracking is incredibly useful during complex operations like groupin import qualified DataFrame as D import qualified DataFrame.Functions as F import qualified DataFrame.Typed as DT +import qualified DataFrame.Typed.Functions as DTF import DataFrame.Expression.Operators ((|>)) @@ -369,7 +374,7 @@ main = do print $ tdf |> DT.groupBy @'["ocean_proximity"] |> DT.aggregate - (DT.as @"total" (DT.count (DT.col @"median_house_value"))) + (DT.as @"total" (DTF.count (DT.col @"median_house_value"))) ``` ## The Trade Off diff --git a/docs/working_with_nulls.md b/docs/working_with_nulls.md index 7cd478f1..f3dbdc95 100644 --- a/docs/working_with_nulls.md +++ b/docs/working_with_nulls.md @@ -79,7 +79,7 @@ These work on non-nullable columns and return plain `Bool`. ```haskell -- Both columns must be non-nullable -D.filterWhere (F.col @Int "id" .==. F.lit 2) df +D.filter (F.col @Int "id" .==. F.lit 2) df ``` ### Nullable-aware operators (`.+`, `.-`, `.*`, `./`, `.==`, `.<`, …) diff --git a/examples/Chipotle.hs b/examples/Chipotle.hs index 5acbecd9..959ae224 100644 --- a/examples/Chipotle.hs +++ b/examples/Chipotle.hs @@ -45,7 +45,7 @@ run = do df |> D.select [F.name item_name, F.name quantity] -- It's more efficient to filter before grouping. - |> D.filterWhere (item_name .==. "Chicken Burrito") + |> D.filter (item_name .==. "Chicken Burrito") |> D.groupBy [F.name item_name] |> D.aggregate [ "sum" .= F.sum quantity @@ -56,7 +56,7 @@ run = do let firstOrder = df - |> D.filterWhere + |> D.filter ( F.lift (maybe False (T.isInfixOf "Guacamole")) choice_description .&&. (item_name .==. "Chicken Bowl") ) diff --git a/examples/OneBillionRowChallenge.hs b/examples/OneBillionRowChallenge.hs index 052e30d8..ebb03ecd 100644 --- a/examples/OneBillionRowChallenge.hs +++ b/examples/OneBillionRowChallenge.hs @@ -26,13 +26,13 @@ run = do readSeparated measurementSchema "../../1brc/measurements.txt" - |> L.groupBy - [F.name measurementNames] + |> L.groupBy [F.name measurementNames] + |> L.aggregate [ F.minimum measurementTemperature `as` "minimum" , F.mean measurementTemperature `as` "mean" , F.maximum measurementTemperature `as` "maximum" ] - |> L.runDataFrame + |> L.run print $ D.sortBy [D.Asc measurementNames] df endCalculation <- getCurrentTime let calculationTime = diffUTCTime endCalculation startCalculation diff --git a/src/DataFrame.hs b/src/DataFrame.hs index bcf1f352..3598a1ce 100644 --- a/src/DataFrame.hs +++ b/src/DataFrame.hs @@ -173,7 +173,7 @@ import DataFrame.Typed.TH.Records as SchemaTH (deriveSchemaFromType, deriveSchem import DataFrame.Lazy as Lazy ( LazyDataFrame, fromDataFrame, - runDataFrame, + run, scanCsv, scanCsvWith, scanParquet, @@ -192,6 +192,7 @@ import DataFrame.Operations.Join as Join ( fullOuterJoin, innerJoin, join, + joinOn, leftJoin, rightJoin, ) @@ -243,7 +244,6 @@ import DataFrame.Operations.Subset as Subset ( filterBy, filterJust, filterNothing, - filterWhere, kFolds, randomSplit, range, diff --git a/src/DataFrame/Typed.hs b/src/DataFrame/Typed.hs index 09f1d050..e70c14f4 100644 --- a/src/DataFrame/Typed.hs +++ b/src/DataFrame/Typed.hs @@ -8,17 +8,13 @@ module DataFrame.Typed ( TypedGrouped, These (..), - -- * Typed expressions + -- * Typed expressions (combinators live in "DataFrame.Typed.Functions") TExpr (..), col, lit, ifThenElse, caseWhen, orElse, - lift, - lift2, - nullLift, - nullLift2, -- * Same-type comparison operators (.==.), @@ -28,11 +24,21 @@ module DataFrame.Typed ( (.>=.), (.>.), + -- * Same-type arithmetic operators + (.+.), + (.-.), + (.*.), + (./.), + (.^^.), + (.^.), + -- * Nullable-aware arithmetic operators (.+), (.-), (.*), (./), + (.^^), + (.^), -- * Nullable-aware comparison operators (three-valued logic) (.==), @@ -45,84 +51,11 @@ module DataFrame.Typed ( -- * Logical operators (.&&.), (.||.), - DataFrame.Typed.Expr.not, - - -- * Aggregation expression combinators - DataFrame.Typed.Expr.sum, - mean, - median, - count, - countAll, - DataFrame.Typed.Expr.minimum, - DataFrame.Typed.Expr.maximum, - collect, - over, - - -- * Expression combinators (full DataFrame.Functions parity) - DataFrame.Typed.Expr.eq, - DataFrame.Typed.Expr.lt, - DataFrame.Typed.Expr.gt, - DataFrame.Typed.Expr.leq, - DataFrame.Typed.Expr.geq, - DataFrame.Typed.Expr.and, - DataFrame.Typed.Expr.or, - DataFrame.Typed.Expr.countDistinct, - DataFrame.Typed.Expr.topK, - DataFrame.Typed.Expr.bottomK, - DataFrame.Typed.Expr.isNull, - DataFrame.Typed.Expr.isNotNull, - DataFrame.Typed.Expr.contains, - DataFrame.Typed.Expr.length, - DataFrame.Typed.Expr.div, - DataFrame.Typed.Expr.mod, - mode, - sumMaybe, - DataFrame.Typed.Expr.meanMaybe, - DataFrame.Typed.Expr.variance, - DataFrame.Typed.Expr.medianMaybe, - DataFrame.Typed.Expr.percentile, - stddev, - stddevMaybe, - zScore, - pow, - relu, - DataFrame.Typed.Expr.min, - DataFrame.Typed.Expr.max, - reduce, - toMaybe, - fromMaybe, - isJust, - isNothing, - fromJust, - whenPresent, - whenBothPresent, - recode, - recodeWithCondition, - recodeWithDefault, - firstOrNothing, - lastOrNothing, - splitOn, - match, - matchAll, - parseDate, - daysBetween, - bind, - - -- * Cast / coercion expressions - DataFrame.Typed.Expr.cast, - DataFrame.Typed.Expr.castWithDefault, - DataFrame.Typed.Expr.castEither, - DataFrame.Typed.Expr.unsafeCast, - castExpr, - castExprWithDefault, - castExprEither, - unsafeCastExpr, - toDouble, + (.&&), + (.||), -- * Typed sort orders TSortOrder (..), - asc, - desc, -- * Freeze / thaw boundary freeze, @@ -142,7 +75,6 @@ module DataFrame.Typed ( toIntMatrix, -- * Schema-preserving operations - filterWhere, filter, filterBy, filterAllJust, @@ -150,6 +82,7 @@ module DataFrame.Typed ( filterNothing, filterAllNothing, sortBy, + topBy, take, takeLast, drop, @@ -162,6 +95,7 @@ module DataFrame.Typed ( -- * Schema-modifying operations derive, + deriveMany, impute, select, exclude, @@ -174,6 +108,12 @@ module DataFrame.Typed ( dropColumn, replaceColumn, + -- * Construction and folds + empty, + null, + fold, + (|>), + -- * Metadata dimensions, nRows, @@ -190,6 +130,9 @@ module DataFrame.Typed ( symmetricDifference, -- * Joins + JoinType (..), + join, + joinOn, innerJoin, leftJoin, rightJoin, @@ -202,7 +145,9 @@ module DataFrame.Typed ( aggregateUntyped, -- * Column transformations - applyColumn, + apply, + applyInt, + applyDouble, applyMany, applyWhere, applyAtIndex, @@ -224,6 +169,24 @@ module DataFrame.Typed ( valueCounts, valueProportions, + -- * Statistics + mean, + meanMaybe, + median, + medianMaybe, + percentile, + genericPercentile, + standardDeviation, + skewness, + variance, + interQuartileRange, + sum, + correlation, + frequencies, + imputeWith, + summarize, + describeColumns, + #ifdef WITH_TH -- * Template Haskell deriveSchema, @@ -242,8 +205,48 @@ module DataFrame.Typed ( -- * Record bridge (ADT <-> TypedDataFrame) HasSchema (..), - fromRecordsTyped, - toRecordsTyped, + fromRecords, + toRecords, + + -- * Display and rendering + DisplayOptions (..), + defaultDisplayOptions, + display, + prettyPrint, + toMarkdown, + toMarkdown', + toCsv, + toCsv', + toSeparated, + + -- * IO + readJSON, + readJSONEither, +#ifdef WITH_CSV + ReadOptions (..), + defaultReadOptions, + readCsv, + readCsvWithError, + readCsvWithOpts, + readTsv, + readSeparated, + writeCsv, + writeTsv, + writeSeparated, +#endif +#ifdef WITH_PARQUET + ParquetReadOptions (..), + ParquetWriteOptions (..), + defaultParquetReadOptions, + defaultParquetWriteOptions, + readParquet, + readParquetWithError, + readParquetWithOpts, + readParquetFiles, + readParquetFilesWithOpts, + writeParquet, + writeParquetWithOptions, +#endif -- * Generics opt-in for schema derivation SchemaOf, @@ -286,7 +289,7 @@ module DataFrame.Typed ( AllKnownSymbol (..), ) where -import Prelude hiding (drop, filter, take) +import Prelude hiding (drop, filter, null, sum, take) import DataFrame.Typed.Access ( columnAsDoubleVector, @@ -306,8 +309,10 @@ import DataFrame.Typed.Aggregate ( groupBy, ) import DataFrame.Typed.Apply ( + apply, + applyDouble, + applyInt, applyAtIndex, - applyColumn, applyMany, applyWhere, deriveWithExpr, @@ -324,7 +329,60 @@ import DataFrame.Typed.Sampling ( stratifiedSample, stratifiedSplit, ) -import DataFrame.Typed.Expr +import DataFrame.Typed.Expr ( + TExpr (..), + caseWhen, + col, + ifThenElse, + lit, + orElse, + (.&&), + (.&&.), + (.*), + (.*.), + (.+), + (.+.), + (.-), + (.-.), + (./), + (./.), + (./=), + (./=.), + (.<), + (.<.), + (.<=), + (.<=.), + (.==), + (.==.), + (.>), + (.>.), + (.>=), + (.>=.), + (.^), + (.^.), + (.^^), + (.^^.), + (.||), + (.||.), + ) +import DataFrame.Typed.Statistics ( + correlation, + describeColumns, + frequencies, + genericPercentile, + imputeWith, + interQuartileRange, + mean, + meanMaybe, + median, + medianMaybe, + percentile, + skewness, + standardDeviation, + summarize, + sum, + variance, + ) import DataFrame.Typed.Freeze (freeze, freezeWithError, thaw, unsafeFreeze) import DataFrame.Typed.Generic ( NameCase (..), @@ -333,13 +391,61 @@ import DataFrame.Typed.Generic ( genericFromColumns, genericToColumns, ) -import DataFrame.Typed.Join (fullOuterJoin, innerJoin, leftJoin, rightJoin) -import DataFrame.Typed.Operations -import DataFrame.Typed.Record ( - HasSchema (..), - fromRecordsTyped, - toRecordsTyped, +import DataFrame.Display (DisplayOptions (..), defaultDisplayOptions) +import DataFrame.Expression.Operators ((|>)) +import DataFrame.Operations.Join (JoinType (..)) +#ifdef WITH_CSV +import DataFrame.IO.CSV (ReadOptions (..), defaultReadOptions) +import DataFrame.Typed.IO.CSV ( + readCsv, + readCsvWithError, + readCsvWithOpts, + readSeparated, + readTsv, + writeCsv, + writeSeparated, + writeTsv, + ) +#endif +#ifdef WITH_PARQUET +import DataFrame.IO.Parquet ( + ParquetReadOptions (..), + ParquetWriteOptions (..), + defaultParquetReadOptions, + defaultParquetWriteOptions, + ) +import DataFrame.Typed.IO.Parquet ( + readParquet, + readParquetFiles, + readParquetFilesWithOpts, + readParquetWithError, + readParquetWithOpts, + writeParquet, + writeParquetWithOptions, + ) +#endif +import DataFrame.Typed.IO.JSON (readJSON, readJSONEither) +import DataFrame.Typed.Interop ( + display, + fromRecords, + prettyPrint, + toCsv, + toCsv', + toMarkdown, + toMarkdown', + toRecords, + toSeparated, ) +import DataFrame.Typed.Join ( + fullOuterJoin, + innerJoin, + join, + joinOn, + leftJoin, + rightJoin, + ) +import DataFrame.Typed.Operations +import DataFrame.Typed.Record (HasSchema (..)) import DataFrame.Typed.Schema #ifdef WITH_TH import DataFrame.Typed.TH ( diff --git a/src/DataFrame/Typed/Interop.hs b/src/DataFrame/Typed/Interop.hs new file mode 100644 index 00000000..0d03df3f --- /dev/null +++ b/src/DataFrame/Typed/Interop.hs @@ -0,0 +1,63 @@ +{-# LANGUAGE ScopedTypeVariables #-} + +-- | Typed display, rendering and record helpers mirroring "DataFrame". +module DataFrame.Typed.Interop ( + -- * Display and rendering + display, + prettyPrint, + toMarkdown, + toMarkdown', + toCsv, + toCsv', + toSeparated, + + -- * Records + fromRecords, + toRecords, +) where + +import qualified Data.Text as T + +import qualified DataFrame.Core as Core +import DataFrame.Display (DisplayOptions) +import qualified DataFrame.Display as Display +import DataFrame.Typed.Freeze (thaw) +import DataFrame.Typed.Record (HasSchema (..), fromRecordsTyped, toRecordsTyped) +import DataFrame.Typed.Types (TExpr (..), TypedDataFrame) + +-- | Print the frame with the given options. +display :: DisplayOptions -> TypedDataFrame cols -> IO () +display opts = Display.display opts . thaw + +-- | Render a typed expression as source-like text. +prettyPrint :: TExpr cols a -> String +prettyPrint (TExpr e) = Core.prettyPrint e + +-- | Render the frame as a Markdown table. +toMarkdown :: TypedDataFrame cols -> T.Text +toMarkdown = Core.toMarkdown . thaw + +-- | 'toMarkdown' as a 'String'. +toMarkdown' :: TypedDataFrame cols -> String +toMarkdown' = Core.toMarkdown' . thaw + +-- | Render the frame as CSV. +toCsv :: TypedDataFrame cols -> T.Text +toCsv = Core.toCsv . thaw + +-- | 'toCsv' as a 'String'. +toCsv' :: TypedDataFrame cols -> String +toCsv' = Core.toCsv' . thaw + +-- | Render the frame with the given separator. +toSeparated :: Char -> TypedDataFrame cols -> T.Text +toSeparated sep = Core.toSeparated sep . thaw + +-- | Build a typed frame from records; the schema comes from the record type. +fromRecords :: forall a. (HasSchema a) => [a] -> TypedDataFrame (Schema a) +fromRecords = fromRecordsTyped + +-- | Parse records back out of a typed frame. +toRecords :: + forall a. (HasSchema a) => TypedDataFrame (Schema a) -> Either T.Text [a] +toRecords = toRecordsTyped diff --git a/tests/LazyParity.hs b/tests/LazyParity.hs index 9291697a..d9976679 100644 --- a/tests/LazyParity.hs +++ b/tests/LazyParity.hs @@ -13,7 +13,6 @@ import qualified DataFrame.Functions as F import qualified DataFrame.IO.CSV as Csv import qualified DataFrame.Internal.Column as DI import qualified DataFrame.Internal.Expression as E -import DataFrame.Lazy (SortOrder (Descending)) import qualified DataFrame.Lazy as L import qualified DataFrame.Operations.Aggregation as Agg import DataFrame.Operations.Join (JoinType (LEFT)) @@ -99,10 +98,11 @@ joinPipelineParity = let customersQ = L.scanCsv customersSchema (T.pack customersPath) lazy <- L.scanCsv ordersSchema (T.pack ordersPath) - |> (\o -> L.join LEFT "customer_id" "customer_id" o customersQ) - |> L.groupBy ["region", "plan"] aggs - |> L.sortBy [("revenue", Descending)] - |> L.runDataFrame + |> (\o -> L.join LEFT ["customer_id"] o customersQ) + |> L.groupBy ["region", "plan"] + |> L.aggregate aggs + |> L.sortBy [L.Desc (E.Col @Double "revenue")] + |> L.run assertEqual "join pipeline lazy == eager" (show eager) @@ -125,9 +125,10 @@ groupByPipelineParity = Agg.groupBy ["customer_id"] ordersDf lazy <- L.scanCsv ordersSchema (T.pack ordersPath) - |> L.groupBy ["customer_id"] aggs - |> L.sortBy [("total", Descending)] - |> L.runDataFrame + |> L.groupBy ["customer_id"] + |> L.aggregate aggs + |> L.sortBy [L.Desc (E.Col @Double "total")] + |> L.run assertEqual "groupBy pipeline lazy == eager" (show eager) @@ -150,9 +151,10 @@ streamingMeanParity = (L.scanCsvStreamingWith Csv.readSeparated ordersSchema (T.pack csvPath)) { L.batchSize = 2 } - |> L.groupBy ["customer_id"] aggs - |> L.sortBy [("customer_id", Descending)] - actual <- L.runDataFrame query + |> L.groupBy ["customer_id"] + |> L.aggregate aggs + |> L.sortBy [L.Desc (E.Col @Int "customer_id")] + actual <- L.run query let expected = Perm.sortBy [Perm.Desc (E.Col @Int "customer_id")] $ Agg.aggregate aggs $ @@ -174,5 +176,45 @@ streamingMeanParity = , ("discount", DI.fromList (replicate 12 (0 :: Double))) ] +-- | Streaming (unbounded) probe side through the streaming hash join. +streamingJoinParity :: Test +streamingJoinParity = + TestCase $ + withCsv (ordersFrame 5_000) $ \ordersPath -> + withCsv (customersFrame 257) $ \customersPath -> do + ordersDf <- Csv.readCsvWithSchema ordersSchema ordersPath + customersDf <- Csv.readCsvWithSchema customersSchema customersPath + let byOrder = Perm.sortBy [Perm.Asc (E.Col @Int "order_id")] + orders = + (L.scanCsvStreaming ordersSchema (T.pack ordersPath)) + { L.batchSize = 512 + } + renamedCustomers = D.rename "customer_id" "cid" customersDf + check lbl eager lazy = do + actual <- L.run lazy + assertEqual lbl (show (byOrder eager)) (show (byOrder actual)) + check + "streaming inner join" + (Join.innerJoin ["customer_id"] ordersDf customersDf) + (L.innerJoin ["customer_id"] orders (L.fromDataFrame customersDf)) + check + "streaming left join" + (Join.leftJoin ["customer_id"] ordersDf customersDf) + (L.leftJoin ["customer_id"] orders (L.fromDataFrame customersDf)) + check + "streaming joinOn" + (Join.joinOn LEFT [("customer_id", "cid")] ordersDf renamedCustomers) + ( L.joinOn + LEFT + [("customer_id", "cid")] + orders + (L.fromDataFrame renamedCustomers) + ) + tests :: [Test] -tests = [joinPipelineParity, groupByPipelineParity, streamingMeanParity] +tests = + [ joinPipelineParity + , groupByPipelineParity + , streamingMeanParity + , streamingJoinParity + ] diff --git a/tests/LazyParquet.hs b/tests/LazyParquet.hs index 7f8fcab2..210dd639 100644 --- a/tests/LazyParquet.hs +++ b/tests/LazyParquet.hs @@ -9,7 +9,6 @@ import qualified Data.Text as T import Data.Time (UTCTime) import qualified DataFrame as D import qualified DataFrame.Functions as F -import DataFrame.Lazy (SortOrder (..)) import qualified DataFrame.Lazy as L import DataFrame.Schema (Schema (..), schemaType) import Test.HUnit @@ -41,7 +40,7 @@ basicScan = TestCase ( do eager <- D.readParquet plainPath - actual <- L.runDataFrame (L.scanParquet allTypesSchema (T.pack plainPath)) + actual <- L.run (L.scanParquet allTypesSchema (T.pack plainPath)) assertEqual "basicScan dimensions" (D.dimensions eager) (D.dimensions actual) ) @@ -51,7 +50,7 @@ columnProjection = TestCase ( do actual <- - L.runDataFrame + L.run (L.select ["id", "bool_col"] (L.scanParquet allTypesSchema (T.pack plainPath))) assertEqual "columnProjection" (8, 2) (D.dimensions actual) ) @@ -62,7 +61,7 @@ filterPushdown = TestCase ( do actual <- - L.runDataFrame + L.run ( L.filter (F.geq (F.col @Int32 "id") (F.lit 6)) (L.scanParquet allTypesSchema (T.pack plainPath)) @@ -76,7 +75,7 @@ filterAndProject = TestCase ( do actual <- - L.runDataFrame + L.run ( L.select ["id", "bool_col"] ( L.filter @@ -93,7 +92,7 @@ multiFileGlob = TestCase ( do actual <- - L.runDataFrame + L.run ( L.scanParquet allTypesSchema "./tests/data/alltypes_plain*.parquet" @@ -109,11 +108,11 @@ sortAndLimit = TestCase ( do actual <- - L.runDataFrame + L.run ( L.take 3 ( L.sortBy - [("id", Ascending)] + [L.Asc (L.col @Int32 "id")] (L.scanParquet allTypesSchema (T.pack plainPath)) ) ) diff --git a/tests/LazyProjection.hs b/tests/LazyProjection.hs index cf1edc82..d62fe544 100644 --- a/tests/LazyProjection.hs +++ b/tests/LazyProjection.hs @@ -43,7 +43,7 @@ twoCols = -- | The scan keeps only the schema's columns, whatever their header position. scanProjects :: Test scanProjects = TestLabel "scan_projects" $ TestCase $ withCsv ',' $ \csvPath -> do - df <- L.runDataFrame (L.scanCsv twoCols (T.pack csvPath)) + df <- L.run (L.scanCsv twoCols (T.pack csvPath)) assertEqual "dims" (2, 2) (D.dimensions df) assertEqual "columns" ["dob", "id"] (D.columnNames df) @@ -52,7 +52,7 @@ scanProjectsWithFastReader :: Test scanProjectsWithFastReader = TestLabel "scan_projects_fast_reader" $ TestCase $ withCsv ',' $ \csvPath -> do df <- - L.runDataFrame + L.run (L.scanCsvWith Fast.fastReadCsvWithOpts twoCols (T.pack csvPath)) assertEqual "dims" (2, 2) (D.dimensions df) assertEqual "columns" ["dob", "id"] (D.columnNames df) @@ -60,9 +60,9 @@ scanProjectsWithFastReader = -- | Both readers must agree on what a scan produces. readersAgree :: Test readersAgree = TestLabel "scan_readers_agree" $ TestCase $ withCsv ',' $ \csvPath -> do - slow <- L.runDataFrame (L.scanCsv twoCols (T.pack csvPath)) + slow <- L.run (L.scanCsv twoCols (T.pack csvPath)) fast <- - L.runDataFrame (L.scanCsvWith Fast.fastReadCsvWithOpts twoCols (T.pack csvPath)) + L.run (L.scanCsvWith Fast.fastReadCsvWithOpts twoCols (T.pack csvPath)) assertEqual "same frame" slow fast -- | scanSeparated must honour its separator, not always split on commas. @@ -70,7 +70,7 @@ scanHonoursSeparator :: Test scanHonoursSeparator = TestLabel "scan_honours_separator" $ TestCase $ withCsv ';' $ \csvPath -> do -- With the separator ignored, the schema's columns are not found. - df <- L.runDataFrame (L.scanSeparated ';' twoCols (T.pack csvPath)) + df <- L.run (L.scanSeparated ';' twoCols (T.pack csvPath)) assertEqual "dims" (2, 2) (D.dimensions df) assertEqual "columns" ["dob", "id"] (D.columnNames df) @@ -80,7 +80,7 @@ scanMissingColumn = TestLabel "scan_missing_column" $ TestCase $ withCsv ',' $ \ let bogus = Schema $ M.fromList [("nope", schemaType @Int)] r <- try - (L.runDataFrame (L.scanCsv bogus (T.pack csvPath)) >>= evaluate . forceDataFrame) + (L.run (L.scanCsv bogus (T.pack csvPath)) >>= evaluate . forceDataFrame) case r of Left (e :: SomeException) -> assertBool diff --git a/tests/Main.hs b/tests/Main.hs index 93ede77f..2e4fb621 100644 --- a/tests/Main.hs +++ b/tests/Main.hs @@ -68,6 +68,7 @@ import qualified Operations.VectorKernel import qualified Operations.Window import qualified Operations.WriteCsv import qualified PackedTextMigration +import qualified Parity.FourWay import qualified Parquet import qualified Plotting import qualified PrettyPrint @@ -139,6 +140,7 @@ tests = ++ Parquet.tests ++ LazyParquet.tests ++ LazyParity.tests + ++ Parity.FourWay.tests ++ LazyProjection.tests ++ Plotting.tests ++ Simplify.tests diff --git a/tests/Operations/Aggregations.hs b/tests/Operations/Aggregations.hs index b3252a87..0418b584 100644 --- a/tests/Operations/Aggregations.hs +++ b/tests/Operations/Aggregations.hs @@ -9,6 +9,7 @@ import qualified DataFrame as D import qualified DataFrame.Functions as F import qualified DataFrame.Internal.Column as DI import qualified DataFrame.Typed as DT +import qualified DataFrame.Typed.Functions as DTF import Data.Function import DataFrame.Expression.Operators @@ -82,8 +83,8 @@ countAllAggregationTyped = , '("test6", Integer) ] & DT.groupBy @'["test1"] - & DT.aggregate (DT.as @"n" DT.countAll) - & DT.sortBy [DT.asc (DT.col @"test1")] + & DT.aggregate (DT.as @"n" DTF.countAll) + & DT.sortBy [DTF.asc (DT.col @"test1")] & DT.thaw ) ) @@ -109,8 +110,8 @@ foldAggregationTyped = , '("test6", Integer) ] & DT.groupBy @'["test1"] - & DT.aggregate (DT.as @"test2_count" (DT.count (DT.col @"test2"))) - & DT.sortBy [DT.asc (DT.col @"test1")] + & DT.aggregate (DT.as @"test2_count" (DTF.count (DT.col @"test2"))) + & DT.sortBy [DTF.asc (DT.col @"test1")] & DT.thaw ) ) @@ -153,8 +154,8 @@ numericAggregationTyped = , '("test6", Integer) ] & DT.groupBy @'["test1"] - & DT.aggregate (DT.as @"test2_mean" (DT.mean (DT.col @"test2"))) - & DT.sortBy [DT.asc (DT.col @"test1")] + & DT.aggregate (DT.as @"test2_mean" (DTF.mean (DT.col @"test2"))) + & DT.sortBy [DTF.asc (DT.col @"test1")] & DT.thaw ) ) diff --git a/tests/Operations/Derive.hs b/tests/Operations/Derive.hs index cbe02715..0683396c 100644 --- a/tests/Operations/Derive.hs +++ b/tests/Operations/Derive.hs @@ -12,6 +12,7 @@ import qualified DataFrame.Functions as F import qualified DataFrame.Internal.Column as DI import qualified DataFrame.Internal.DataFrame as DI import qualified DataFrame.Typed as DT +import qualified DataFrame.Typed.Functions as DTF import Test.HUnit @@ -56,10 +57,10 @@ deriveWAITyped = ( DT.columnAsList @"test4" $ DT.derive @"test4" - ( DT.lift2 + ( DTF.lift2 (++) - (DT.lift show (DT.col @"test1")) - (DT.lift (: ([] :: [Char])) (DT.col @"test3")) + (DTF.lift show (DT.col @"test1")) + (DTF.lift (: ([] :: [Char])) (DT.col @"test3")) ) ( either (error . show) diff --git a/tests/Operations/Filter.hs b/tests/Operations/Filter.hs index f8ede1b9..c9e0c1c3 100644 --- a/tests/Operations/Filter.hs +++ b/tests/Operations/Filter.hs @@ -5,6 +5,7 @@ module Operations.Filter where import qualified Data.Text as T import qualified DataFrame as D +import DataFrame.Expression.Operators ((.==.)) import qualified DataFrame.Functions as F import qualified DataFrame.Internal.Column as DI @@ -32,8 +33,8 @@ filterColumnDoesNotExist = TestCase ( assertExpectException "[Error Case]" - (D.columnNotFound "test0" "filter" (D.columnNames testData)) - (print $ D.filter (F.col @Int "test0") even testData) + (D.columnNotFound "test0" "" (D.columnNames testData)) + (print $ D.filter (F.col @Int "test0" .==. 0) testData) ) filterColumnWrongType :: Test @@ -42,7 +43,7 @@ filterColumnWrongType = ( assertExpectException "[Error Case]" (D.typeMismatchError (show $ typeRep @Integer) (show $ typeRep @Int)) - (print $ D.filter (F.col @Integer "test1") even testData) + (print $ D.filter (F.col @Integer "test1" .==. 0) testData) ) filterByColumnDoesNotExist :: Test @@ -50,7 +51,7 @@ filterByColumnDoesNotExist = TestCase ( assertExpectException "[Error Case]" - (D.columnNotFound "test0" "filter" (D.columnNames testData)) + (D.columnNotFound "test0" "filterBy" (D.columnNames testData)) (print $ D.filterBy even (F.col @Int "test0") testData) ) @@ -69,7 +70,7 @@ filterColumnInexistentValues = ( assertEqual "Non existent filter value returns no rows" (0, 8) - (D.dimensions $ D.filter (F.col @Int "test1") (< 0) testData) + (D.dimensions $ D.filterBy (< 0) (F.col @Int "test1") testData) ) filterColumnAllValues :: Test @@ -78,7 +79,7 @@ filterColumnAllValues = ( assertEqual "Filters all columns" (26, 8) - (D.dimensions $ D.filter (F.col @Int "test1") (const True) testData) + (D.dimensions $ D.filterBy (const True) (F.col @Int "test1") testData) ) filterJustWAI :: Test diff --git a/tests/Operations/Join.hs b/tests/Operations/Join.hs index e827dd63..b904fa1d 100644 --- a/tests/Operations/Join.hs +++ b/tests/Operations/Join.hs @@ -12,6 +12,7 @@ import qualified DataFrame.Functions as F import DataFrame.Internal.Column.Types (These (..)) import DataFrame.Operations.Join import qualified DataFrame.Typed as DT +import qualified DataFrame.Typed.Functions as DTF import Test.HUnit df1 :: D.DataFrame @@ -102,7 +103,7 @@ testInnerJoinTyped = , ("B", D.fromList ["B0" :: Text, "B1", "B2"]) ] ) - (DT.thaw $ DT.sortBy [DT.asc (DT.col @"key")] (DT.innerJoin @'["key"] tdf1 tdf2)) + (DT.thaw $ DT.sortBy [DTF.asc (DT.col @"key")] (DT.innerJoin @'["key"] tdf1 tdf2)) ) testLeftJoinTyped :: Test @@ -116,7 +117,7 @@ testLeftJoinTyped = , ("B", D.fromList [Just "B0", Just "B1" :: Maybe Text, Just "B2"]) ] ) - (DT.thaw $ DT.sortBy [DT.asc (DT.col @"key")] (DT.leftJoin @'["key"] tdf1 tdf2)) + (DT.thaw $ DT.sortBy [DTF.asc (DT.col @"key")] (DT.leftJoin @'["key"] tdf1 tdf2)) ) -- A right-hand frame whose payload column is already optional. @@ -147,7 +148,7 @@ testLeftJoinTypedOptional = ) ] ) - (DT.thaw $ DT.sortBy [DT.asc (DT.col @"key")] joined) + (DT.thaw $ DT.sortBy [DTF.asc (DT.col @"key")] joined) ) where joined :: @@ -166,7 +167,7 @@ testRightJoinTyped = , ("B", D.fromList ["B0" :: Text, "B1", "B2"]) ] ) - (DT.thaw $ DT.sortBy [DT.asc (DT.col @"key")] (DT.rightJoin @'["key"] tdf1 tdf2)) + (DT.thaw $ DT.sortBy [DTF.asc (DT.col @"key")] (DT.rightJoin @'["key"] tdf1 tdf2)) ) staffDf :: D.DataFrame diff --git a/tests/Operations/OperatorIdentity.hs b/tests/Operations/OperatorIdentity.hs index 625eb689..13742615 100644 --- a/tests/Operations/OperatorIdentity.hs +++ b/tests/Operations/OperatorIdentity.hs @@ -76,7 +76,7 @@ powVariants = two = F.lit (2 :: Int) rowsKept :: D.DataFrame -> Expr Bool -> Int -rowsKept df p = D.nRows (D.filterWhere p df) +rowsKept df p = D.nRows (D.filter p df) -- Both sides must keep the same rows before and after simplification. sameRows :: String -> D.DataFrame -> Expr Bool -> Test @@ -204,7 +204,10 @@ everyFunctionIsBuiltin = , ("fromJust", hasKey (F.fromJust m)) , ("recode", hasKey (F.recode [(1 :: Int, 'a')] x)) , ("recodeWithDefault", hasKey (F.recodeWithDefault 'z' [(1 :: Int, 'a')] x)) - , ("recodeWithCondition", hasKey (F.recodeWithCondition (F.lit (0 :: Int)) [((.> F.lit @Int 1), 1)] x)) + , + ( "recodeWithCondition" + , hasKey (F.recodeWithCondition (F.lit (0 :: Int)) [((.> F.lit @Int 1), 1)] x) + ) , ("firstOrNothing", hasKey (F.firstOrNothing l)) , ("lastOrNothing", hasKey (F.lastOrNothing l)) , ("contains", hasKey (F.contains "a" t)) diff --git a/tests/Operations/ParallelGroupBy.hs b/tests/Operations/ParallelGroupBy.hs index 44460cd9..d6037b81 100644 --- a/tests/Operations/ParallelGroupBy.hs +++ b/tests/Operations/ParallelGroupBy.hs @@ -222,7 +222,7 @@ skewedTextParity filtered = TestCase $ do | otherwise = T.pack ("k" ++ show ((i * 7919) `mod` 40000)) df0 = D.fromColumns [("k", packedCol (map key [0 .. n - 1])), mkInt "v" [0 .. n - 1]] - df = if filtered then D.filterWhere (F.col @T.Text "k" ./=. "") df0 else df0 + df = if filtered then D.filter (F.col @T.Text "k" ./=. "") df0 else df0 assertBool "skewed packed text: par == seq" (sameGrouping (groupByPar ["k"] df) (groupBySeq ["k"] df)) diff --git a/tests/Operations/Record.hs b/tests/Operations/Record.hs index 26d29e75..4fb298a5 100644 --- a/tests/Operations/Record.hs +++ b/tests/Operations/Record.hs @@ -125,8 +125,8 @@ orderSample = basicTypedRoundTrip :: Test basicTypedRoundTrip = TestCase $ do let df :: DT.TypedDataFrame OrderSchema - df = DT.fromRecordsTyped orderSample - case DT.toRecordsTyped df of + df = DT.fromRecords orderSample + case DT.toRecords df of Left e -> assertFailure (T.unpack e) Right xs -> assertEqual "typed round-trip" orderSample xs @@ -173,8 +173,8 @@ schemaNameOverride :: Test schemaNameOverride = TestCase $ do let xs = [Item 1 9.99, Item 2 19.99] df :: DT.TypedDataFrame ItemsSchema - df = DT.fromRecordsTyped xs - case DT.toRecordsTyped df of + df = DT.fromRecords xs + case DT.toRecords df of Left e -> assertFailure (T.unpack e) Right ys -> assertEqual "schema-name override round-trip" xs ys @@ -323,7 +323,7 @@ deriveSchemaAccessorFilter = TestCase $ do , Order 4 "us" 25.0 ] big = - D.filterWhere + D.filter (orderAmount .>. F.lit @Double 15.0 .&&. orderRegion .==. F.lit @T.Text "us") df case D.toRecords big :: Either T.Text [Order] of @@ -350,9 +350,9 @@ deriveSchemaAccessorDerive = TestCase $ do labelColumnFilter :: Test labelColumnFilter = TestCase $ do let df :: DT.TypedDataFrame OrderSchema - df = DT.fromRecordsTyped orderSample - usOnly = DT.filterWhere (#region DT..==. "us") df - case DT.toRecordsTyped usOnly of + df = DT.fromRecords orderSample + usOnly = DT.filter (#region DT..==. "us") df + case DT.toRecords usOnly of Left e -> assertFailure (T.unpack e) Right xs -> assertEqual diff --git a/tests/Operations/Window.hs b/tests/Operations/Window.hs index 73d313a8..bbd44fb0 100644 --- a/tests/Operations/Window.hs +++ b/tests/Operations/Window.hs @@ -33,7 +33,7 @@ globalFilterTest = "Global median filter removes global outliers" 7 ( purchases - & D.filterWhere + & D.filter (F.col @Double "amount" .<=. F.median (F.col @Double "amount") * 10) & D.nRows ) @@ -50,7 +50,7 @@ overMedianFilterTest = where actual = purchases - & D.filterWhere + & D.filter ( F.col @Double "amount" .<=. F.over ["country"] (F.median (F.col @Double "amount")) * 10 @@ -76,13 +76,13 @@ globalVsOverDifferentResults = where globalResult = purchases - & D.filterWhere + & D.filter (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 + & D.filter ( F.col @Double "amount" .<=. F.over ["country"] (F.median (F.col @Double "amount")) * 3 @@ -198,7 +198,7 @@ blogPostExampleGlobal = discount = F.col @Double "discount" actual = purchases - & D.filterWhere (amount .<=. F.median amount * 10) + & D.filter (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")] @@ -222,7 +222,7 @@ blogPostExampleOver = discount = F.col @Double "discount" actual = purchases - & D.filterWhere (amount .<=. F.over ["country"] (F.median amount) * 10) + & D.filter (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")] diff --git a/tests/Parity/FourWay.hs b/tests/Parity/FourWay.hs new file mode 100644 index 00000000..2234c605 --- /dev/null +++ b/tests/Parity/FourWay.hs @@ -0,0 +1,206 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +-- | Every operation, run through eager, typed, lazy and typed lazy APIs. +module Parity.FourWay (tests) 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.Lazy as L +import qualified DataFrame.Typed as DT +import qualified DataFrame.Typed.Functions as DTF +import qualified DataFrame.Typed.Lazy as TL + +import Test.HUnit + +type Orders = + '[ '("cid", Int) + , '("region", T.Text) + , '("amount", Double) + ] + +type Customers = + '[ '("cid", Int) + , '("region", T.Text) + , '("tier", T.Text) + ] + +type Users = '[ '("uid", Int), '("name", T.Text)] + +ordersDF :: D.DataFrame +ordersDF = + D.fromColumns + [ ("cid", DI.fromList [1, 2, 2, 3, 5, 1 :: Int]) + , ("region", DI.fromList ["eu", "us", "us", "eu", "us", "us" :: T.Text]) + , ("amount", DI.fromList [10, 20, 5, 7.5, 3, 1 :: Double]) + ] + +customersDF :: D.DataFrame +customersDF = + D.fromColumns + [ ("cid", DI.fromList [1, 2, 3, 4 :: Int]) + , ("region", DI.fromList ["eu", "us", "eu", "eu" :: T.Text]) + , ("tier", DI.fromList ["gold", "free", "free", "gold" :: T.Text]) + ] + +usersDF :: D.DataFrame +usersDF = + D.fromColumns + [ ("uid", DI.fromList [1, 2, 4 :: Int]) + , ("name", DI.fromList ["ann", "bo", "cy" :: T.Text]) + ] + +ordersT :: DT.TypedDataFrame Orders +ordersT = either (error . show) id (DT.freezeWithError ordersDF) + +customersT :: DT.TypedDataFrame Customers +customersT = either (error . show) id (DT.freezeWithError customersDF) + +usersT :: DT.TypedDataFrame Users +usersT = either (error . show) id (DT.freezeWithError usersDF) + +-- | Compare all four results against the eager one. +fourWay :: + String -> + D.DataFrame -> + DT.TypedDataFrame cols -> + L.LazyDataFrame -> + TL.TypedLazyDataFrame cols -> + Test +fourWay lbl eager typed lazy typedLazy = TestCase $ do + assertEqual (lbl ++ ": typed == eager") eager (DT.thaw typed) + lazyDf <- L.run lazy + assertEqual (lbl ++ ": lazy == eager") eager lazyDf + typedLazyDf <- TL.run typedLazy + assertEqual (lbl ++ ": typed lazy == eager") eager (DT.thaw typedLazyDf) + +lo, lc, lu :: L.LazyDataFrame +lo = L.fromDataFrame ordersDF +lc = L.fromDataFrame customersDF +lu = L.fromDataFrame usersDF + +tlo :: TL.TypedLazyDataFrame Orders +tlo = TL.fromDataFrame ordersT + +tlc :: TL.TypedLazyDataFrame Customers +tlc = TL.fromDataFrame customersT + +tlu :: TL.TypedLazyDataFrame Users +tlu = TL.fromDataFrame usersT + +joins :: Test +joins = + TestList + [ fourWay + "innerJoin multi-key" + (D.innerJoin ["cid", "region"] ordersDF customersDF) + (DT.innerJoin @'["cid", "region"] ordersT customersT) + (L.innerJoin ["cid", "region"] lo lc) + (TL.innerJoin @'["cid", "region"] tlo tlc) + , fourWay + "join INNER" + (D.join D.INNER ["cid"] ordersDF customersDF) + (DT.join @'DT.INNER @'["cid"] ordersT customersT) + (L.join L.INNER ["cid"] lo lc) + (TL.join @'TL.INNER @'["cid"] tlo tlc) + , fourWay + "leftJoin" + (D.leftJoin ["cid"] ordersDF customersDF) + (DT.leftJoin @'["cid"] ordersT customersT) + (L.leftJoin ["cid"] lo lc) + (TL.leftJoin @'["cid"] tlo tlc) + , fourWay + "rightJoin" + (D.rightJoin ["cid"] ordersDF customersDF) + (DT.rightJoin @'["cid"] ordersT customersT) + (L.rightJoin ["cid"] lo lc) + (TL.rightJoin @'["cid"] tlo tlc) + , fourWay + "fullOuterJoin" + (D.fullOuterJoin ["cid"] ordersDF customersDF) + (DT.fullOuterJoin @'["cid"] ordersT customersT) + (L.fullOuterJoin ["cid"] lo lc) + (TL.fullOuterJoin @'["cid"] tlo tlc) + , fourWay + "joinOn differing key names" + (D.joinOn D.LEFT [("cid", "uid")] ordersDF usersDF) + (DT.joinOn @'DT.LEFT @'[ '("cid", "uid")] ordersT usersT) + (L.joinOn L.LEFT [("cid", "uid")] lo lu) + (TL.joinOn @'TL.LEFT @'[ '("cid", "uid")] tlo tlu) + ] + +groupAndSort :: Test +groupAndSort = + TestList + [ fourWay + "groupBy + aggregate + sortBy" + ( ordersDF + D.|> D.groupBy ["region"] + D.|> D.aggregate + [ F.sum (F.col @Double "amount") `D.as` "total" + , F.count (F.col @Double "amount") `D.as` "n" + ] + D.|> D.sortBy [D.Asc (F.col @T.Text "region")] + ) + ( ordersT + DT.|> DT.groupBy @'["region"] + DT.|> DT.aggregate + ( DT.as @"n" (DTF.count (DT.col @"amount")) + . DT.as @"total" (DTF.sum (DT.col @"amount")) + ) + DT.|> DT.sortBy [DT.Asc (DT.col @"region")] + ) + ( lo + L.|> L.groupBy ["region"] + L.|> L.aggregate + [ F.sum (F.col @Double "amount") `L.as` "total" + , F.count (F.col @Double "amount") `L.as` "n" + ] + L.|> L.sortBy [L.Asc (L.col @T.Text "region")] + ) + ( tlo + TL.|> TL.groupBy @'["region"] + TL.|> TL.aggregate + ( TL.as @"n" (DTF.count (TL.col @"amount")) + . TL.as @"total" (DTF.sum (TL.col @"amount")) + ) + TL.|> TL.sortBy [TL.Asc (TL.col @"region")] + ) + , fourWay + "sortBy expression, filter, derive, select, take" + ( ordersDF + D.|> D.filter (F.col @Double "amount" D..>. 2) + D.|> D.derive "twice" (F.col @Double "amount" * 2) + D.|> D.sortBy [D.Desc (F.col @Double "amount" * (-1))] + D.|> D.select ["cid", "twice"] + D.|> D.take 3 + ) + ( ordersT + DT.|> DT.filter (DT.col @"amount" DT..>. 2) + DT.|> DT.derive @"twice" (DT.col @"amount" * 2) + DT.|> DT.sortBy [DT.Desc (DT.col @"amount" * (-1))] + DT.|> DT.select @'["cid", "twice"] + DT.|> DT.take 3 + ) + ( lo + L.|> L.filter (F.col @Double "amount" L..>. 2) + L.|> L.derive "twice" (F.col @Double "amount" * 2) + L.|> L.sortBy [L.Desc (F.col @Double "amount" * (-1))] + L.|> L.select ["cid", "twice"] + L.|> L.take 3 + ) + ( tlo + TL.|> TL.filter (TL.col @"amount" TL..>. 2) + TL.|> TL.derive @"twice" (TL.col @"amount" * 2) + TL.|> TL.sortBy [TL.Desc (TL.col @"amount" * (-1))] + TL.|> TL.select @'["cid", "twice"] + TL.|> TL.take 3 + ) + ] + +tests :: [Test] +tests = [joins, groupAndSort] diff --git a/tests/Parquet.hs b/tests/Parquet.hs index 49f8699e..df92dc5f 100644 --- a/tests/Parquet.hs +++ b/tests/Parquet.hs @@ -99,7 +99,7 @@ allTypesPlainSnappy = testBothReadParquetPaths $ \readParquet -> TestCase ( assertEqual "allTypesPlainSnappy" - (D.filter (F.col @Int32 "id") (`elem` [6, 7]) allTypes) + (D.filterBy (`elem` [6, 7]) (F.col @Int32 "id") allTypes) (unsafePerformIO (readParquet "./tests/data/alltypes_plain.snappy.parquet")) ) @@ -108,7 +108,7 @@ allTypesDictionary = testBothReadParquetPaths $ \readParquet -> TestCase ( assertEqual "allTypesPlainSnappy" - (D.filter (F.col @Int32 "id") (`elem` [0, 1]) allTypes) + (D.filterBy (`elem` [0, 1]) (F.col @Int32 "id") allTypes) (unsafePerformIO (readParquet "./tests/data/alltypes_dictionary.parquet")) ) diff --git a/tests/Typed/Parity.hs b/tests/Typed/Parity.hs index 5d23a2de..d58b11b1 100644 --- a/tests/Typed/Parity.hs +++ b/tests/Typed/Parity.hs @@ -10,10 +10,12 @@ import qualified Data.Text as T import System.Random (mkStdGen) import qualified DataFrame as D +import DataFrame.Expression.Operators ((./=.)) import qualified DataFrame.Functions as F import qualified DataFrame.Internal.Column as DI import DataFrame.Internal.Expression (eqExpr) import qualified DataFrame.Typed as DT +import qualified DataFrame.Typed.Functions as DTF import qualified DataFrame.Typed.Statistics as TS import Test.HUnit @@ -75,39 +77,39 @@ statParity = exprParity :: Test exprParity = TestList - [ deriveEq "pow" (DT.pow (DT.col @"x") 2) (F.pow (F.col @Int "x") 2) - , deriveEq "relu" (DT.relu (DT.col @"x")) (F.relu (F.col @Int "x")) - , deriveEq "toMaybe" (DT.toMaybe (DT.col @"x")) (F.toMaybe (F.col @Int "x")) + [ deriveEq "pow" (DTF.pow (DT.col @"x") 2) (F.pow (F.col @Int "x") 2) + , deriveEq "relu" (DTF.relu (DT.col @"x")) (F.relu (F.col @Int "x")) + , deriveEq "toMaybe" (DTF.toMaybe (DT.col @"x")) (F.toMaybe (F.col @Int "x")) , deriveEq "min" - (DT.min (DT.col @"x") (DT.lit 3)) + (DTF.min (DT.col @"x") (DT.lit 3)) (F.min (F.col @Int "x") (F.lit @Int 3)) , deriveEq "splitOn" - (DT.splitOn "a" (DT.col @"g")) + (DTF.splitOn "a" (DT.col @"g")) (F.splitOn "a" (F.col @T.Text "g")) , deriveEq "recodeWithDefault" - (DT.recodeWithDefault 0 [(1 :: Int, 100 :: Int)] (DT.col @"x")) + (DTF.recodeWithDefault 0 [(1 :: Int, 100 :: Int)] (DT.col @"x")) (F.recodeWithDefault 0 [(1 :: Int, 100 :: Int)] (F.col @Int "x")) , deriveEq "eq" - (DT.eq (DT.col @"x") (DT.lit 3)) + (DTF.eq (DT.col @"x") (DT.lit 3)) (F.eq (F.col @Int "x") (F.lit @Int 3)) , deriveEq "and" - (DT.and (DT.lt (DT.col @"x") (DT.lit 5)) (DT.geq (DT.col @"x") (DT.lit 2))) + (DTF.and (DTF.lt (DT.col @"x") (DT.lit 5)) (DTF.geq (DT.col @"x") (DT.lit 2))) (F.and (F.lt x (F.lit @Int 5)) (F.geq x (F.lit @Int 2))) - , deriveEq "contains" (DT.contains "a" (DT.col @"g")) (F.contains "a" g) - , deriveEq "length" (DT.length (DT.col @"g")) (F.length g) + , deriveEq "contains" (DTF.contains "a" (DT.col @"g")) (F.contains "a" g) + , deriveEq "length" (DTF.length (DT.col @"g")) (F.length g) , deriveEq "isNull" - (DT.isNull (DT.toMaybe (DT.col @"x"))) + (DTF.isNull (DTF.toMaybe (DT.col @"x"))) (F.isNull (F.toMaybe x)) - , deriveEq "cast" (DT.cast @"x" @Int) (F.cast @Int "x") + , deriveEq "cast" (DTF.cast @"x" @Int) (F.cast @Int "x") , deriveEq "castWithDefault" - (DT.castWithDefault @"g" (0 :: Int)) + (DTF.castWithDefault @"g" (0 :: Int)) (F.castWithDefault (0 :: Int) "g") ] where @@ -128,10 +130,10 @@ exprParity = aggParity :: Test aggParity = TestList - [ same "countDistinct" (DT.countDistinct (DT.col @"x")) (F.countDistinct x) - , same "topK" (DT.topK 2 (DT.col @"x")) (F.topK 2 x) - , same "bottomK" (DT.bottomK 2 (DT.col @"x")) (F.bottomK 2 x) - , same "not" (DT.not (DT.eq (DT.col @"x") 1)) (F.not (F.eq x 1)) + [ same "countDistinct" (DTF.countDistinct (DT.col @"x")) (F.countDistinct x) + , same "topK" (DTF.topK 2 (DT.col @"x")) (F.topK 2 x) + , same "bottomK" (DTF.bottomK 2 (DT.col @"x")) (F.bottomK 2 x) + , same "not" (DTF.not (DTF.eq (DT.col @"x") 1)) (F.not (F.eq x 1)) ] where x = F.col @Int "x" @@ -144,9 +146,9 @@ applyParity = TestList [ TestCase ( assertEqual - "applyColumn type-change" + "apply type-change" (D.apply (show :: Int -> String) "x" baseDF) - (DT.thaw (DT.applyColumn @"x" (show :: Int -> String) tdf)) + (DT.thaw (DT.apply @"x" (show :: Int -> String) tdf)) ) , TestCase ( assertEqual @@ -231,9 +233,73 @@ numericDF = numericT :: DT.TypedDataFrame '[ '("x", Int), '("y", Double)] numericT = either (error . show) id (DT.freezeWithError numericDF) +-- Ops added for import-swap parity: same result through either API. +newOpsParity :: Test +newOpsParity = + TestList + [ same + "topBy" + (D.topBy 1 2 [D.Desc (F.col @Double "y")] baseDF) + (DT.topBy 1 2 [DT.Desc (DT.col @"y")] tdf) + , same + "deriveMany" + ( D.deriveMany + [ "x2" D..= (F.col @Int "x" * 2) + , "y2" D..= (F.col @Double "y" + 1) + ] + baseDF + ) + ( DT.deriveMany + (DT.as @"y2" (DT.col @"y" + 1) . DT.as @"x2" (DT.col @"x" * 2)) + tdf + ) + , same + "join INNER" + (D.join D.INNER ["x"] leftDF numericDF) + (DT.join @'DT.INNER @'["x"] leftT numericT) + , same + "join LEFT" + (D.join D.LEFT ["x"] numericDF leftDF) + (DT.join @'DT.LEFT @'["x"] numericT leftT) + , same + "applyInt" + (D.applyInt (* 10) "x" baseDF) + (DT.applyInt @"x" (* 10) tdf) + , same + "applyDouble" + (D.applyDouble (> 20) "y" baseDF) + (DT.applyDouble @"y" (> 20) tdf) + , same + "fold" + (D.fold (\n -> D.filter (F.col @Int "x" ./=. F.lit n)) [2, 4] baseDF) + (DT.fold (\n -> DT.filter (DT.col @"x" DT../=. DT.lit n)) [2, 4] tdf) + , TestCase (assertBool "empty is null" (DT.null DT.empty)) + , TestCase (assertEqual "null" (D.null baseDF) (DT.null tdf)) + , TestCase (assertEqual "toCsv" (D.toCsv baseDF) (DT.toCsv tdf)) + , TestCase + (assertEqual "toMarkdown" (D.toMarkdown baseDF) (DT.toMarkdown tdf)) + , TestCase + ( assertEqual + "frame stats via DataFrame.Typed" + (D.mean (F.col @Double "y") baseDF) + (DT.mean (DT.col @"y") tdf) + ) + , TestCase + ( assertEqual + "readJSONEither" + (D.toCsv <$> D.readJSONEither jsonBytes) + (DT.toCsv <$> DT.readJSONEither @'[ '("a", Maybe Double)] jsonBytes) + ) + ] + where + same :: String -> D.DataFrame -> DT.TypedDataFrame cols -> Test + same lbl untyped typed = TestCase (assertEqual lbl untyped (DT.thaw typed)) + jsonBytes = "[{\"a\": 1}, {\"a\": 2}]" + tests :: [Test] tests = - [ statParity + [ newOpsParity + , statParity , exprParity , aggParity , applyParity From 628086406f8ff4064f8dba451a5adc7b6072a155 Mon Sep 17 00:00:00 2001 From: Michael Chavinda Date: Thu, 22 Oct 2026 20:00:56 -0700 Subject: [PATCH 2/3] feat: full lazy and typed-lazy op coverage with four-way parity tests --- CHANGELOG.md | 13 +- README.md | 17 + dataframe-lazy/dataframe-lazy.cabal | 3 + dataframe-lazy/src/DataFrame/Lazy.hs | 5 + .../src/DataFrame/Lazy/Internal/Collect.hs | 255 ++++++ .../src/DataFrame/Lazy/Internal/DataFrame.hs | 347 ++++++++ .../src/DataFrame/Lazy/Internal/Executor.hs | 96 ++- .../DataFrame/Lazy/Internal/LogicalPlan.hs | 33 + .../src/DataFrame/Lazy/Internal/Optimizer.hs | 70 +- .../DataFrame/Lazy/Internal/PhysicalPlan.hs | 7 +- dataframe-lazy/src/DataFrame/Typed/Lazy.hs | 762 +++++++++++++++++- docs/base_scripts/base_readme.md | 17 + tests/LazyParity.hs | 71 +- tests/LazyParquet.hs | 13 +- tests/LazyProjection.hs | 8 +- tests/Parity/FourWay.hs | 456 ++++++++++- 16 files changed, 2131 insertions(+), 42 deletions(-) create mode 100644 dataframe-lazy/src/DataFrame/Lazy/Internal/Collect.hs diff --git a/CHANGELOG.md b/CHANGELOG.md index 04bb8abe..6278fff2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -50,6 +50,15 @@ `DataFrame.Internal.Statistics`. ### Additions +* `DataFrame.Lazy` and `DataFrame.Typed.Lazy` cover the eager API: row ops + (`filterBy`, `filterJust`, `drop`, `takeLast`, `range`, `topBy`, + `distinct`, `sample`, `shuffle`, splits and folds), column ops (`exclude`, + `rename`, `deriveMany`, `impute`, `apply*`, `insert*`, `cloneColumn`, …), + merges and set operations (`<>`/`append`, `|||`, `union`, `intersect`, + `difference`, `symmetricDifference`), summaries, and value-returning ops + (`mean`, `nRows`, `columnAs*`, `toCsv`, `writeCsv`, `writeParquet`, + `display`, …) that run the query and return `IO`. Row-local ops stream + batch by batch; ops that need the whole input collect it first. * `joinOn jt [(leftKey, rightKey)]` in all four APIs (typed: `joinOn @'INNER @'[ '("id", "uid")]`); the result keeps the left names. * Typed parity with the untyped API: `topBy`, `deriveMany` (built with @@ -65,8 +74,8 @@ `contains` and `length`, completing parity with `DataFrame.Functions`. ### Fixes -* Streaming CSV scans return columns in file order, like eager reads, - instead of alphabetical order. +* Lazy CSV and Parquet scans return columns in file order, like eager + reads, instead of alphabetical order. * A user function named like a builtin (e.g. `"add"`) no longer runs the builtin's kernel instead of the user's function. * `eqExpr` no longer equates different `lift`/`lift2` functions, `fromMaybe` diff --git a/README.md b/README.md index 7da8a85d..d97b3840 100644 --- a/README.md +++ b/README.md @@ -514,6 +514,23 @@ D.take 10 lazyResult |> D.toMarkdown' The optimizer pushes the filter into the scan, drops unreferenced columns before reading, and stops pulling batches once 1000 rows have been collected. +### Switching between APIs + +The eager, typed, lazy and typed-lazy APIs share operation names and argument +order, so moving a pipeline between them is mostly a change of imports: + +| | Frame ops (`D`) | Expressions (`F`) | +|---|---|---| +| Eager | `DataFrame` | `DataFrame.Functions` | +| Typed | `DataFrame.Typed` | `DataFrame.Typed.Functions` | +| Lazy | `DataFrame.Lazy` | `DataFrame.Functions` | +| Typed lazy | `DataFrame.Typed.Lazy` | `DataFrame.Typed.Functions` | + +A lazy pipeline starts from `scanCsv`/`scanParquet`/`fromDataFrame` and ends with +`run`; ops that return a value (`mean`, `nRows`, `writeCsv`, …) run the query +and return `IO`. Typed APIs name columns at the type level (`col @"x"`, +`select @'["x"]`) where untyped ones take text. + ## Documentation * User guide: https://dataframe.readthedocs.io/en/latest/ diff --git a/dataframe-lazy/dataframe-lazy.cabal b/dataframe-lazy/dataframe-lazy.cabal index 50d4653f..58b3f95b 100644 --- a/dataframe-lazy/dataframe-lazy.cabal +++ b/dataframe-lazy/dataframe-lazy.cabal @@ -36,6 +36,7 @@ library -- sealed. SortOrder(..)/LazyDataFrame re-exported via DataFrame.Lazy. other-modules: DataFrame.Lazy.Internal.DataFrame + DataFrame.Lazy.Internal.Collect DataFrame.Lazy.Internal.Executor DataFrame.Lazy.Internal.LogicalPlan DataFrame.Lazy.Internal.Optimizer @@ -50,9 +51,11 @@ library dataframe-operations >= 2.6 && < 2.7, dataframe-parquet >= 1.6 && < 1.7, dataframe-parsing >= 2.2 && < 2.3, + dataframe-viz >= 1.3.2 && < 1.4, directory >= 1.3.0.0 && < 2, filepath >= 1.4 && < 2, Glob >= 0.10 && < 1, + random >= 1 && < 2, stm >= 2.5 && < 3, temporary >= 1.3 && < 2, text >= 2.1 && < 3, diff --git a/dataframe-lazy/src/DataFrame/Lazy.hs b/dataframe-lazy/src/DataFrame/Lazy.hs index 759df414..27072621 100644 --- a/dataframe-lazy/src/DataFrame/Lazy.hs +++ b/dataframe-lazy/src/DataFrame/Lazy.hs @@ -1,12 +1,17 @@ -- | Lazy, streaming queries with the same operations as "DataFrame". module DataFrame.Lazy ( module DataFrame.Lazy.Internal.DataFrame, + module DataFrame.Lazy.Internal.Collect, module DataFrame.Expression.Operators, SortOrder (..), JoinType (..), + DisplayOptions (..), + defaultDisplayOptions, ) where +import DataFrame.Display (DisplayOptions (..), defaultDisplayOptions) import DataFrame.Expression.Operators +import DataFrame.Lazy.Internal.Collect import DataFrame.Lazy.Internal.DataFrame import DataFrame.Operations.Join (JoinType (..)) import DataFrame.Operations.Permutation (SortOrder (..)) diff --git a/dataframe-lazy/src/DataFrame/Lazy/Internal/Collect.hs b/dataframe-lazy/src/DataFrame/Lazy/Internal/Collect.hs new file mode 100644 index 00000000..443aadac --- /dev/null +++ b/dataframe-lazy/src/DataFrame/Lazy/Internal/Collect.hs @@ -0,0 +1,255 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ScopedTypeVariables #-} + +-- | Lazy ops that return values: each runs the query, then the eager op. +module DataFrame.Lazy.Internal.Collect ( + -- * Shape + dimensions, + nRows, + nColumns, + columnNames, + null, + + -- * Statistics + mean, + meanMaybe, + median, + medianMaybe, + percentile, + genericPercentile, + standardDeviation, + skewness, + variance, + interQuartileRange, + sum, + correlation, + valueCounts, + valueProportions, + + -- * Column and row access + columnAsVector, + columnAsList, + columnAsIntVector, + columnAsDoubleVector, + columnAsFloatVector, + columnAsUnboxedVector, + toDoubleMatrix, + toFloatMatrix, + toIntMatrix, + toRowList, + toRowVector, + toRecords, + safeApply, + + -- * Rendering and output + display, + toCsv, + toCsv', + toMarkdown, + toMarkdown', + toSeparated, + writeCsv, + writeTsv, + writeSeparated, + writeParquet, + writeParquetWithOptions, +) where + +import qualified Data.Text as T +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU +import Prelude hiding (null, sum) + +import DataFrame.Display (DisplayOptions) +import qualified DataFrame.Display as Display +import DataFrame.Errors (DataFrameException) +import qualified DataFrame.IO.CSV as CSV +import DataFrame.IO.Parquet (ParquetWriteOptions) +import qualified DataFrame.IO.Parquet as Parquet +import qualified DataFrame.Internal.Column as C +import qualified DataFrame.Internal.DataFrame as D +import qualified DataFrame.Internal.Expression as E +import DataFrame.Internal.Row (Any, Row) +import qualified DataFrame.Internal.Row as Row +import DataFrame.Lazy.Internal.DataFrame (LazyDataFrame, run, select) +import qualified DataFrame.Operations.Core as Core +import qualified DataFrame.Operations.Statistics as Stats +import qualified DataFrame.Operations.Transformations as Trans +import DataFrame.Typed.Record (HasSchema) +import qualified DataFrame.Typed.Record as Record + +-- | Run only the columns an expression reads. +runFor :: E.Expr a -> LazyDataFrame -> IO D.DataFrame +runFor e ldf = case E.getColumns e of + [] -> run ldf + cols -> run (select cols ldf) + +-- | Rows and columns of the result. +dimensions :: LazyDataFrame -> IO (Int, Int) +dimensions ldf = Core.dimensions <$> run ldf + +-- | Number of rows in the result. +nRows :: LazyDataFrame -> IO Int +nRows ldf = Core.nRows <$> run ldf + +-- | Number of columns in the result. +nColumns :: LazyDataFrame -> IO Int +nColumns ldf = Core.nColumns <$> run ldf + +-- | Column names of the result. +columnNames :: LazyDataFrame -> IO [T.Text] +columnNames ldf = D.columnNames <$> run ldf + +-- | True when the result has no rows or no columns. +null :: LazyDataFrame -> IO Bool +null ldf = D.null <$> run ldf + +mean :: + (C.Columnable a, Real a, VU.Unbox a) => E.Expr a -> LazyDataFrame -> IO Double +mean e ldf = Stats.mean e <$> runFor e ldf + +meanMaybe :: + (C.Columnable a, Real a) => E.Expr (Maybe a) -> LazyDataFrame -> IO Double +meanMaybe e ldf = Stats.meanMaybe e <$> runFor e ldf + +median :: + (C.Columnable a, Real a, VU.Unbox a) => E.Expr a -> LazyDataFrame -> IO Double +median e ldf = Stats.median e <$> runFor e ldf + +medianMaybe :: + (C.Columnable a, Real a) => E.Expr (Maybe a) -> LazyDataFrame -> IO Double +medianMaybe e ldf = Stats.medianMaybe e <$> runFor e ldf + +percentile :: + (C.Columnable a, Real a, VU.Unbox a) => + Int -> E.Expr a -> LazyDataFrame -> IO Double +percentile p e ldf = Stats.percentile p e <$> runFor e ldf + +genericPercentile :: + (C.Columnable a, Ord a) => Int -> E.Expr a -> LazyDataFrame -> IO a +genericPercentile p e ldf = Stats.genericPercentile p e <$> runFor e ldf + +standardDeviation :: + (C.Columnable a, Real a, VU.Unbox a) => E.Expr a -> LazyDataFrame -> IO Double +standardDeviation e ldf = Stats.standardDeviation e <$> runFor e ldf + +skewness :: + (C.Columnable a, Real a, VU.Unbox a) => E.Expr a -> LazyDataFrame -> IO Double +skewness e ldf = Stats.skewness e <$> runFor e ldf + +variance :: + (C.Columnable a, Real a, VU.Unbox a) => E.Expr a -> LazyDataFrame -> IO Double +variance e ldf = Stats.variance e <$> runFor e ldf + +interQuartileRange :: + (C.Columnable a, Real a, VU.Unbox a) => E.Expr a -> LazyDataFrame -> IO Double +interQuartileRange e ldf = Stats.interQuartileRange e <$> runFor e ldf + +sum :: (C.Columnable a, Num a) => E.Expr a -> LazyDataFrame -> IO a +sum e ldf = Stats.sum e <$> runFor e ldf + +-- | Pearson correlation of two named columns. +correlation :: T.Text -> T.Text -> LazyDataFrame -> IO (Maybe Double) +correlation a b ldf = Stats.correlation a b <$> run (select [a, b] ldf) + +valueCounts :: + (Ord a, C.Columnable a) => E.Expr a -> LazyDataFrame -> IO [(a, Int)] +valueCounts e ldf = Core.valueCounts e <$> runFor e ldf + +valueProportions :: + (Ord a, C.Columnable a) => E.Expr a -> LazyDataFrame -> IO [(a, Double)] +valueProportions e ldf = Core.valueProportions e <$> runFor e ldf + +columnAsVector :: + (C.Columnable a) => + E.Expr a -> LazyDataFrame -> IO (Either DataFrameException (V.Vector a)) +columnAsVector e ldf = Core.columnAsVector e <$> runFor e ldf + +columnAsList :: (C.Columnable a) => E.Expr a -> LazyDataFrame -> IO [a] +columnAsList e ldf = Core.columnAsList e <$> runFor e ldf + +columnAsIntVector :: + (C.Columnable a, Num a) => + E.Expr a -> LazyDataFrame -> IO (Either DataFrameException (VU.Vector Int)) +columnAsIntVector e ldf = Core.columnAsIntVector e <$> runFor e ldf + +columnAsDoubleVector :: + (C.Columnable a, Num a) => + E.Expr a -> LazyDataFrame -> IO (Either DataFrameException (VU.Vector Double)) +columnAsDoubleVector e ldf = Core.columnAsDoubleVector e <$> runFor e ldf + +columnAsFloatVector :: + (C.Columnable a, Num a) => + E.Expr a -> LazyDataFrame -> IO (Either DataFrameException (VU.Vector Float)) +columnAsFloatVector e ldf = Core.columnAsFloatVector e <$> runFor e ldf + +columnAsUnboxedVector :: + (C.Columnable a, VU.Unbox a) => + E.Expr a -> LazyDataFrame -> IO (Either DataFrameException (VU.Vector a)) +columnAsUnboxedVector e ldf = Core.columnAsUnboxedVector e <$> runFor e ldf + +toDoubleMatrix :: + LazyDataFrame -> IO (Either DataFrameException (V.Vector (VU.Vector Double))) +toDoubleMatrix ldf = Core.toDoubleMatrix <$> run ldf + +toFloatMatrix :: + LazyDataFrame -> IO (Either DataFrameException (V.Vector (VU.Vector Float))) +toFloatMatrix ldf = Core.toFloatMatrix <$> run ldf + +toIntMatrix :: + LazyDataFrame -> IO (Either DataFrameException (V.Vector (VU.Vector Int))) +toIntMatrix ldf = Core.toIntMatrix <$> run ldf + +toRowList :: LazyDataFrame -> IO [[(T.Text, Any)]] +toRowList ldf = Row.toRowList <$> run ldf + +toRowVector :: [T.Text] -> LazyDataFrame -> IO (V.Vector Row) +toRowVector names ldf = Row.toRowVector names <$> run ldf + +-- | Parse records out of the result. +toRecords :: (HasSchema a) => LazyDataFrame -> IO (Either T.Text [a]) +toRecords ldf = Record.toRecords <$> run ldf + +-- | Run the query, then 'apply' without throwing. +safeApply :: + (C.Columnable b, C.Columnable c) => + (b -> c) -> + T.Text -> + LazyDataFrame -> + IO (Either DataFrameException D.DataFrame) +safeApply f name ldf = Trans.safeApply f name <$> run ldf + +display :: DisplayOptions -> LazyDataFrame -> IO () +display opts ldf = run ldf >>= Display.display opts + +toCsv :: LazyDataFrame -> IO T.Text +toCsv ldf = D.toCsv <$> run ldf + +toCsv' :: LazyDataFrame -> IO String +toCsv' ldf = D.toCsv' <$> run ldf + +toMarkdown :: LazyDataFrame -> IO T.Text +toMarkdown ldf = D.toMarkdown <$> run ldf + +toMarkdown' :: LazyDataFrame -> IO String +toMarkdown' ldf = D.toMarkdown' <$> run ldf + +toSeparated :: Char -> LazyDataFrame -> IO T.Text +toSeparated sep ldf = D.toSeparated sep <$> run ldf + +writeCsv :: FilePath -> LazyDataFrame -> IO () +writeCsv path ldf = run ldf >>= CSV.writeCsv path + +writeTsv :: FilePath -> LazyDataFrame -> IO () +writeTsv path ldf = run ldf >>= CSV.writeTsv path + +writeSeparated :: Char -> FilePath -> LazyDataFrame -> IO () +writeSeparated sep path ldf = run ldf >>= CSV.writeSeparated sep path + +writeParquet :: FilePath -> LazyDataFrame -> IO () +writeParquet path ldf = run ldf >>= Parquet.writeParquet path + +writeParquetWithOptions :: + ParquetWriteOptions -> FilePath -> LazyDataFrame -> IO () +writeParquetWithOptions opts path ldf = + run ldf >>= Parquet.writeParquetWithOptions opts path diff --git a/dataframe-lazy/src/DataFrame/Lazy/Internal/DataFrame.hs b/dataframe-lazy/src/DataFrame/Lazy/Internal/DataFrame.hs index 1b97d7e0..aacdce03 100644 --- a/dataframe-lazy/src/DataFrame/Lazy/Internal/DataFrame.hs +++ b/dataframe-lazy/src/DataFrame/Lazy/Internal/DataFrame.hs @@ -8,6 +8,8 @@ module DataFrame.Lazy.Internal.DataFrame where import qualified Data.Text as T +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU import DataFrame.IO.CSV ( CsvBytesReader, CsvReader, @@ -17,15 +19,33 @@ import DataFrame.IO.CSV ( import qualified DataFrame.Internal.Column as C import qualified DataFrame.Internal.DataFrame as D import qualified DataFrame.Internal.Expression as E +import DataFrame.Internal.Expression.Operators.Nullable (BaseType) import DataFrame.Lazy.Internal.Executor (execute) import DataFrame.Lazy.Internal.LogicalPlan ( + BatchFn (..), + BatchFn2 (..), + ColumnEffect (..), DataSource (..), LogicalPlan (..), + opaque, ) import qualified DataFrame.Lazy.Internal.Optimizer as Opt +import qualified DataFrame.Operations.Aggregation as Agg +import qualified DataFrame.Operations.Core as Core import DataFrame.Operations.Join (JoinType (..)) +import DataFrame.Operations.Merge () +import qualified DataFrame.Operations.Merge as Merge import DataFrame.Operations.Permutation (SortOrder) +import qualified DataFrame.Operations.Permutation as Perm +import qualified DataFrame.Operations.SetOps as SetOps +import qualified DataFrame.Operations.Statistics as Stats +import DataFrame.Operations.Subset (SelectionCriteria, SplittableGen) +import qualified DataFrame.Operations.Subset as Sub +import DataFrame.Operations.Transformations (ImputeOp) +import qualified DataFrame.Operations.Transformations as Trans import DataFrame.Schema (Schema) +import System.Random (RandomGen) +import Prelude hiding (drop, filter, take) -- | A lazy query: a 'LogicalPlan' whose execution waits for 'run'. data LazyDataFrame = LazyDataFrame @@ -166,3 +186,330 @@ sortBy ords ldf = ldf{plan = Sort ords (plan ldf)} -- | Retain at most @n@ rows. take :: Int -> LazyDataFrame -> LazyDataFrame take n ldf = ldf{plan = Limit n (plan ldf)} + +-- * Generic nodes + +-- | Run a row-local eager op on each batch as it streams. +mapBatch :: + T.Text -> + ColumnEffect -> + (D.DataFrame -> D.DataFrame) -> + LazyDataFrame -> + LazyDataFrame +mapBatch lbl eff f ldf = ldf{plan = MapBatch lbl eff (BatchFn f) (plan ldf)} + +-- | Collect the whole input, then run an eager op on it. +barrier :: + T.Text -> + ColumnEffect -> + (D.DataFrame -> D.DataFrame) -> + LazyDataFrame -> + LazyDataFrame +barrier lbl eff f ldf = ldf{plan = Barrier lbl eff (BatchFn f) (plan ldf)} + +-- | Collect both inputs, then run an eager two-frame op on them. +barrier2 :: + T.Text -> + (D.DataFrame -> D.DataFrame -> D.DataFrame) -> + LazyDataFrame -> + LazyDataFrame -> + LazyDataFrame +barrier2 lbl f l r = l{plan = Barrier2 lbl (BatchFn2 f) (plan l) (plan r)} + +-- | Effect of an op that reads the given columns and defines others. +effect :: [T.Text] -> [T.Text] -> ColumnEffect +effect rs = ColumnEffect (Just rs) + +-- | Effect of an op that only looks at row positions, never at values. +positional :: ColumnEffect +positional = effect [] [] + +-- | Concatenate rows, as the eager @<>@ does. +instance Semigroup LazyDataFrame where + (<>) = barrier2 "append" (<>) + +instance Monoid LazyDataFrame where + mempty = fromDataFrame D.empty + +-- * Row operations + +-- | Keep rows whose value of the expression satisfies the predicate. +filterBy :: + (C.Columnable a) => (a -> Bool) -> E.Expr a -> LazyDataFrame -> LazyDataFrame +filterBy p e = mapBatch "filterBy" (effect (E.getColumns e) []) (Sub.filterBy p e) + +-- | Drop rows with Nothing in the column and unwrap its Justs. +filterJust :: T.Text -> LazyDataFrame -> LazyDataFrame +filterJust c = mapBatch "filterJust" (effect [c] [c]) (Sub.filterJust c) + +-- | Keep only rows with Nothing in the column. +filterNothing :: T.Text -> LazyDataFrame -> LazyDataFrame +filterNothing c = mapBatch "filterNothing" (effect [c] []) (Sub.filterNothing c) + +-- | Keep rows with no Nothing in any column. +filterAllJust :: LazyDataFrame -> LazyDataFrame +filterAllJust = mapBatch "filterAllJust" opaque Sub.filterAllJust + +-- | Keep rows that are Nothing in every column. +filterAllNothing :: LazyDataFrame -> LazyDataFrame +filterAllNothing = mapBatch "filterAllNothing" opaque Sub.filterAllNothing + +-- | Skip the first @n@ rows. +drop :: Int -> LazyDataFrame -> LazyDataFrame +drop n ldf = ldf{plan = Offset (max 0 n) (plan ldf)} + +-- | Keep the last @n@ rows. +takeLast :: Int -> LazyDataFrame -> LazyDataFrame +takeLast n = barrier "takeLast" positional (Sub.takeLast n) + +-- | Drop the last @n@ rows. +dropLast :: Int -> LazyDataFrame -> LazyDataFrame +dropLast n = barrier "dropLast" positional (Sub.dropLast n) + +-- | Rows @[start, end)@, clamped like the eager @range@. +range :: (Int, Int) -> LazyDataFrame -> LazyDataFrame +range (start, end) = take (max 0 (end - start')) . drop start' + where + start' = max 0 start + +-- | The first @rows@ rows of the first @cols@ columns. +cube :: (Int, Int) -> LazyDataFrame -> LazyDataFrame +cube (rows, cols) = + mapBatch "cube" opaque (Sub.selectBy [Sub.byIndexRange (0, cols - 1)]) + . take rows + +-- | Rows at the given positions. +selectRows :: [Int] -> LazyDataFrame -> LazyDataFrame +selectRows ixs = barrier "selectRows" positional (Sub.selectRows ixs) + +-- | Keep each row with probability @p@. +sample :: (RandomGen g) => g -> Double -> LazyDataFrame -> LazyDataFrame +sample g p = barrier "sample" positional (Sub.sample g p) + +-- | Shuffle the rows. +shuffle :: (RandomGen g) => g -> LazyDataFrame -> LazyDataFrame +shuffle g = barrier "shuffle" positional (Perm.shuffle g) + +-- | Drop duplicate rows. +distinct :: LazyDataFrame -> LazyDataFrame +distinct = barrier "distinct" opaque Agg.distinct + +-- | Rows @[offset, offset + count)@ after sorting by the given orders. +topBy :: Int -> Int -> [SortOrder] -> LazyDataFrame -> LazyDataFrame +topBy offset count ords = + barrier + "topBy" + (effect (concatMap Perm.sortOrderColumns ords) []) + (Perm.topBy offset count ords) + +-- | Split into a sample of fraction @p@ and the remainder. +randomSplit :: + (RandomGen g) => g -> Double -> LazyDataFrame -> (LazyDataFrame, LazyDataFrame) +randomSplit g p ldf = + ( barrier "randomSplit" positional (fst . Sub.randomSplit g p) ldf + , barrier "randomSplit" positional (snd . Sub.randomSplit g p) ldf + ) + +-- | Split into @k@ folds. +kFolds :: (RandomGen g) => g -> Int -> LazyDataFrame -> [LazyDataFrame] +kFolds g k ldf = + [ barrier "kFolds" positional ((!! i) . Sub.kFolds g k) ldf + | i <- [0 .. k - 1] + ] + +-- | Sample a fraction of each stratum of the expression. +stratifiedSample :: + (SplittableGen g, C.Columnable a) => + g -> Double -> E.Expr a -> LazyDataFrame -> LazyDataFrame +stratifiedSample g p e = + barrier + "stratifiedSample" + (effect (E.getColumns e) []) + (Sub.stratifiedSample g p e) + +-- | Split each stratum of the expression into a sample and the remainder. +stratifiedSplit :: + (SplittableGen g, C.Columnable a) => + g -> Double -> E.Expr a -> LazyDataFrame -> (LazyDataFrame, LazyDataFrame) +stratifiedSplit g p e ldf = + ( barrier "stratifiedSplit" eff (fst . Sub.stratifiedSplit g p e) ldf + , barrier "stratifiedSplit" eff (snd . Sub.stratifiedSplit g p e) ldf + ) + where + eff = effect (E.getColumns e) [] + +-- * Column operations + +-- | Add several columns, each computed from the input frame. +deriveMany :: [E.NamedExpr] -> LazyDataFrame -> LazyDataFrame +deriveMany nes = + mapBatch + "deriveMany" + (effect (concatMap (\(_, E.SomeExpr e) -> E.getColumns e) nes) (map fst nes)) + (Trans.deriveMany nes) + +-- | 'derive', also returning a reference to the new column. +deriveWithExpr :: + (C.Columnable a) => + T.Text -> E.Expr a -> LazyDataFrame -> (E.Expr a, LazyDataFrame) +deriveWithExpr name e ldf = (E.Col name, derive name e ldf) + +-- | Select columns by name, property or range. +selectBy :: [SelectionCriteria] -> LazyDataFrame -> LazyDataFrame +selectBy crit = mapBatch "selectBy" opaque (Sub.selectBy crit) + +-- | Drop the listed columns. +exclude :: [T.Text] -> LazyDataFrame -> LazyDataFrame +exclude cols = mapBatch "exclude" (effect cols []) (Sub.exclude cols) + +-- | Rename a column. +rename :: T.Text -> T.Text -> LazyDataFrame -> LazyDataFrame +rename old new = renameMany [(old, new)] + +-- | Rename columns by @(old, new)@ pairs. +renameMany :: [(T.Text, T.Text)] -> LazyDataFrame -> LazyDataFrame +renameMany pairs ldf = ldf{plan = Rename pairs (plan ldf)} + +-- | Copy a column under a new name. +cloneColumn :: T.Text -> T.Text -> LazyDataFrame -> LazyDataFrame +cloneColumn old new = + mapBatch "cloneColumn" (effect [old] [new]) (Core.cloneColumn old new) + +-- | Insert a column from a container aligned with the whole result. +insert :: + (C.Columnable a, Foldable t) => T.Text -> t a -> LazyDataFrame -> LazyDataFrame +insert name xs = barrier "insert" (effect [] [name]) (Core.insert name xs) + +-- | Insert a boxed vector aligned with the whole result. +insertVector :: + (C.Columnable a) => T.Text -> V.Vector a -> LazyDataFrame -> LazyDataFrame +insertVector name xs = + barrier "insertVector" (effect [] [name]) (Core.insertVector name xs) + +-- | Insert an unboxed vector aligned with the whole result. +insertUnboxedVector :: + (C.Columnable a, VU.Unbox a) => + T.Text -> VU.Vector a -> LazyDataFrame -> LazyDataFrame +insertUnboxedVector name xs = + barrier + "insertUnboxedVector" + (effect [] [name]) + (Core.insertUnboxedVector name xs) + +-- | Insert a container, padding with the default to the result's length. +insertWithDefault :: + (C.Columnable a, Foldable t) => + a -> T.Text -> t a -> LazyDataFrame -> LazyDataFrame +insertWithDefault d name xs = + barrier + "insertWithDefault" + (effect [] [name]) + (Core.insertWithDefault d name xs) + +-- | Insert a vector, padding with the default to the result's length. +insertVectorWithDefault :: + (C.Columnable a) => a -> T.Text -> V.Vector a -> LazyDataFrame -> LazyDataFrame +insertVectorWithDefault d name xs = + barrier + "insertVectorWithDefault" + (effect [] [name]) + (Core.insertVectorWithDefault d name xs) + +-- | Insert a prepared column aligned with the whole result. +insertColumn :: T.Text -> C.Column -> LazyDataFrame -> LazyDataFrame +insertColumn name c = barrier "insertColumn" (effect [] [name]) (D.insertColumn name c) + +-- | Replace Nothing in a column with a value. +impute :: + (ImputeOp a) => E.Expr a -> BaseType a -> LazyDataFrame -> LazyDataFrame +impute e v = mapBatch "impute" (effect cols cols) (Trans.impute e v) + where + cols = E.getColumns e + +-- | Replace Nothing with a statistic of the whole column. +imputeWith :: + (ImputeOp a, C.Columnable (BaseType a)) => + (E.Expr (BaseType a) -> E.Expr (BaseType a)) -> + E.Expr a -> + LazyDataFrame -> + LazyDataFrame +imputeWith f e = barrier "imputeWith" (effect cols cols) (Stats.imputeWith f e) + where + cols = E.getColumns e + +-- | Map a function over a column. +apply :: + (C.Columnable b, C.Columnable c) => + (b -> c) -> T.Text -> LazyDataFrame -> LazyDataFrame +apply f name = mapBatch "apply" (effect [name] [name]) (Trans.apply f name) + +-- | 'apply' over an Int column. +applyInt :: + (C.Columnable b) => (Int -> b) -> T.Text -> LazyDataFrame -> LazyDataFrame +applyInt = apply + +-- | 'apply' over a Double column. +applyDouble :: + (C.Columnable b) => (Double -> b) -> T.Text -> LazyDataFrame -> LazyDataFrame +applyDouble = apply + +-- | Map a function over several columns. +applyMany :: + (C.Columnable b, C.Columnable c) => + (b -> c) -> [T.Text] -> LazyDataFrame -> LazyDataFrame +applyMany f names = mapBatch "applyMany" (effect names names) (Trans.applyMany f names) + +-- | Map over the target column where the criterion column passes. +applyWhere :: + (C.Columnable a, C.Columnable b) => + (a -> Bool) -> T.Text -> (b -> b) -> T.Text -> LazyDataFrame -> LazyDataFrame +applyWhere p fc f tc = + mapBatch "applyWhere" (effect [fc, tc] [tc]) (Trans.applyWhere p fc f tc) + +-- | Map a function over one row of a column. +applyAtIndex :: + (C.Columnable a) => Int -> (a -> a) -> T.Text -> LazyDataFrame -> LazyDataFrame +applyAtIndex i f name = + barrier "applyAtIndex" (effect [name] [name]) (Trans.applyAtIndex i f name) + +-- | Thread the query through a step for each element. +fold :: + (a -> LazyDataFrame -> LazyDataFrame) -> [a] -> LazyDataFrame -> LazyDataFrame +fold f xs ldf = foldl (flip f) ldf xs + +-- * Merges and set operations + +-- | Place two results side by side. +(|||) :: LazyDataFrame -> LazyDataFrame -> LazyDataFrame +(|||) = barrier2 "|||" (Merge.|||) + +-- | Rows in either result, without duplicates. +union :: LazyDataFrame -> LazyDataFrame -> LazyDataFrame +union = barrier2 "union" SetOps.union + +-- | Rows in both results. +intersect :: LazyDataFrame -> LazyDataFrame -> LazyDataFrame +intersect = barrier2 "intersect" SetOps.intersect + +-- | Rows in the first result but not the second. +difference :: LazyDataFrame -> LazyDataFrame -> LazyDataFrame +difference = barrier2 "difference" SetOps.difference + +-- | Rows in exactly one of the results. +symmetricDifference :: LazyDataFrame -> LazyDataFrame -> LazyDataFrame +symmetricDifference = barrier2 "symmetricDifference" SetOps.symmetricDifference + +-- * Frame-valued summaries + +-- | Summary statistics per numeric column. +summarize :: LazyDataFrame -> LazyDataFrame +summarize = barrier "summarize" opaque Stats.summarize + +-- | Name, type and null counts per column. +describeColumns :: LazyDataFrame -> LazyDataFrame +describeColumns = barrier "describeColumns" opaque Core.describeColumns + +-- | Value counts and proportions of an expression. +frequencies :: + (C.Columnable a, Ord a) => E.Expr a -> LazyDataFrame -> LazyDataFrame +frequencies e = barrier "frequencies" opaque (Stats.frequencies e) diff --git a/dataframe-lazy/src/DataFrame/Lazy/Internal/Executor.hs b/dataframe-lazy/src/DataFrame/Lazy/Internal/Executor.hs index e6437996..f4554d27 100644 --- a/dataframe-lazy/src/DataFrame/Lazy/Internal/Executor.hs +++ b/dataframe-lazy/src/DataFrame/Lazy/Internal/Executor.hs @@ -25,6 +25,7 @@ import qualified Data.ByteString as BS import qualified Data.ByteString.Char8 as C8 import qualified Data.ByteString.Unsafe as BSU import Data.IORef +import qualified Data.List as L import qualified Data.Map as M import qualified Data.Maybe import qualified Data.Set as S @@ -39,11 +40,17 @@ import DataFrame.IO.CSV ( schemaReadOptions, ) import qualified DataFrame.IO.Parquet as Parquet +import DataFrame.IO.Parquet.Thrift (unField) +import qualified DataFrame.IO.Parquet.Thrift as Thrift import qualified DataFrame.Internal.Column as C import qualified DataFrame.Internal.DataFrame as D import qualified DataFrame.Internal.Expression as E import qualified DataFrame.Lazy.IO.Binary as Bin -import DataFrame.Lazy.Internal.LogicalPlan (DataSource (..)) +import DataFrame.Lazy.Internal.LogicalPlan ( + BatchFn (..), + BatchFn2 (..), + DataSource (..), + ) import DataFrame.Lazy.Internal.PhysicalPlan import qualified DataFrame.Operations.Aggregation as Agg import qualified DataFrame.Operations.Core as Core @@ -72,6 +79,10 @@ materialized df = do writeIORef ref Nothing return mb +-- | Apply a function to every batch of a stream. +mapStream :: (D.DataFrame -> D.DataFrame) -> Stream -> Stream +mapStream f s = Stream (fmap f <$> pullBatch s) + -- | Drain a stream and concatenate its batches into one DataFrame. collectStream :: Stream -> IO D.DataFrame collectStream stream = go [] @@ -106,6 +117,11 @@ isBounded (PhysicalSort _ c) = isBounded c isBounded (PhysicalHashAggregate _ _ c) = isBounded c isBounded (PhysicalSpill c _) = isBounded c isBounded (PhysicalSourceDF _ _) = True +isBounded (PhysicalRename _ c) = isBounded c +isBounded (PhysicalOffset _ c) = isBounded c +isBounded (PhysicalMapBatch _ _ c) = isBounded c +isBounded (PhysicalBarrier _ _ c) = isBounded c +isBounded (PhysicalBarrier2 _ _ l r) = isBounded l && isBounded r isBounded (PhysicalHashJoin _ _ l r) = isBounded l && isBounded r isBounded (PhysicalSortMergeJoin _ _ l r) = isBounded l && isBounded r @@ -162,6 +178,32 @@ buildStream (PhysicalDerive name uexpr child) = do mb <- pullBatch childStream return $ fmap (Trans.deriveMany [(name, uexpr)]) mb ) +buildStream (PhysicalRename pairs child) = + mapStream (Core.renameMany pairs) <$> buildStream child +buildStream (PhysicalMapBatch _ (BatchFn f) child) = + mapStream f <$> buildStream child +buildStream (PhysicalBarrier _ (BatchFn f) child) = + execute child >>= materialized . f +buildStream (PhysicalBarrier2 _ (BatchFn2 f) l r) = do + leftDf <- execute l + rightDf <- execute r + materialized (f leftDf rightDf) +buildStream (PhysicalOffset n child) = do + childStream <- buildStream child + skippedRef <- newIORef (0 :: Int) + let next = do + mb <- pullBatch childStream + case mb of + Nothing -> return Nothing + Just df -> do + skipped <- readIORef skippedRef + let rows = Core.nRows df + toSkip = min rows (n - skipped) + writeIORef skippedRef (skipped + toSkip) + if toSkip == rows && rows > 0 + then next + else return (Just (Sub.drop toSkip df)) + return (Stream next) buildStream (PhysicalLimit n child) = do childStream <- buildStream child countRef <- newIORef (0 :: Int) @@ -461,7 +503,12 @@ executeParquetScan path cfg = do [] -> return Nothing (f : rest) -> do writeIORef ref rest - Just <$> Parquet.readParquetWithOpts opts f + -- One footer parse orders the selection and feeds the read. + metadata <- Parquet.readMetadataFromPath f + let names = Parquet.getColumnNames (drop 1 (unField (Thrift.schema metadata))) + opts' = + opts{Parquet.selectedColumns = fileOrder names <$> Parquet.selectedColumns opts} + Just <$> Parquet.readParquetWithMetadata Nothing metadata opts' f -- | Read options for a CSV scan: schema gives types, source the separator. scanReadOptions :: Char -> ScanConfig -> ReadOptions @@ -472,9 +519,10 @@ scanReadOptions sep cfg = executeCsvScan :: FilePath -> Char -> CsvReader -> ScanConfig -> IO Stream executeCsvScan path sep reader cfg = do nCaps <- getNumCapabilities + header <- withFile path ReadMode C8.hGetLine chunkPaths <- splitCsvAtNewlines (max 1 nCaps) path - let opts = scanReadOptions sep cfg + let opts = inFileOrder (headerNames sep header) (scanReadOptions sep cfg) batchSz = scanBatchSize cfg chunkDfs <- mapConcurrently (reader opts) chunkPaths mapM_ removeFile chunkPaths @@ -497,6 +545,31 @@ executeCsvScan path sep reader cfg = do in return (Just df') ) +-- | Column names in a CSV header line. +headerNames :: Char -> BS.ByteString -> [T.Text] +headerNames sep = + map (T.dropAround (\c -> c == '"' || c == '\r') . T.pack . C8.unpack) + . C8.split sep + . stripBom + where + stripBom b = Data.Maybe.fromMaybe b (BS.stripPrefix "\xEF\xBB\xBF" b) + +-- | Read the selected columns in file order, like eager reads do. +inFileOrder :: [T.Text] -> ReadOptions -> ReadOptions +inFileOrder names opts = opts{readColumns = fileOrder names <$> readColumns opts} + +-- | Order requested names as they appear in a file (matching by leaf name). +fileOrder :: [T.Text] -> [T.Text] -> [T.Text] +fileOrder names requested = + L.sortOn position requested + where + leaf = T.takeWhileEnd (/= '.') + -- Unknown names trail, so the reader still reports them as missing. + position r = + Data.Maybe.fromMaybe + (length names) + (L.findIndex (\n -> n == r || leaf n == r) names) + executeCsvScanStreaming :: FilePath -> Char -> CsvReader -> ScanConfig -> IO Stream executeCsvScanStreaming path sep reader = @@ -523,11 +596,7 @@ executeCsvScanStreamingWith path sep reader cfg = do ( withFile path ReadMode $ \h -> do rawHeader <- C8.hGetLine h let header = stripUtf8Bom rawHeader - opts = - schemaOpts - { readColumns = - fileOrder (headerNames header) <$> readColumns schemaOpts - } + opts = inFileOrder (headerNames sep header) schemaOpts let feed bytes = unless (BS.null bytes) $ do let window = header <> BS.singleton nl <> bytes @@ -575,17 +644,6 @@ executeCsvScanStreamingWith path sep reader cfg = do stripUtf8Bom bytes = Data.Maybe.fromMaybe bytes (BS.stripPrefix "\xEF\xBB\xBF" bytes) - headerNames = - map (T.dropAround (\c -> c == '"' || c == '\r') . T.pack . C8.unpack) - . C8.split sep - - -- File order, like eager reads; unknown names trail so they still error. - fileOrder names requested = - let wanted = S.fromList requested - present = S.fromList names - in filter (`S.member` wanted) names - <> filter (`S.notMember` present) requested - -- | Find the last record-ending LF in a CSV buffer. lastCompleteRecordNewline :: BS.ByteString -> Maybe Int lastCompleteRecordNewline bytes diff --git a/dataframe-lazy/src/DataFrame/Lazy/Internal/LogicalPlan.hs b/dataframe-lazy/src/DataFrame/Lazy/Internal/LogicalPlan.hs index 906fb022..790bb146 100644 --- a/dataframe-lazy/src/DataFrame/Lazy/Internal/LogicalPlan.hs +++ b/dataframe-lazy/src/DataFrame/Lazy/Internal/LogicalPlan.hs @@ -27,6 +27,29 @@ instance Show DataSource where "CsvSourceStreamingBytes " ++ show path ++ " " ++ show sep ++ " " show (ParquetSource path) = "ParquetSource " ++ show path +-- | Columns a generic node reads and defines; 'Nothing' reads means all. +data ColumnEffect = ColumnEffect + { effectReads :: Maybe [T.Text] + , effectWrites :: [T.Text] + } + deriving (Show) + +-- | An effect the optimizer cannot see through. +opaque :: ColumnEffect +opaque = ColumnEffect Nothing [] + +-- | An eager frame function carried by a generic node. +newtype BatchFn = BatchFn (D.DataFrame -> D.DataFrame) + +instance Show BatchFn where + show _ = "" + +-- | An eager two-frame function carried by a 'Barrier2' node. +newtype BatchFn2 = BatchFn2 (D.DataFrame -> D.DataFrame -> D.DataFrame) + +instance Show BatchFn2 where + show _ = "" + -- | Relational-algebra tree describing what the query computes. data LogicalPlan = -- | Read columns described by the schema from a source. @@ -45,6 +68,16 @@ data LogicalPlan Sort [SortOrder] LogicalPlan | -- | Retain at most N rows. Limit Int LogicalPlan + | -- | Rename columns by @(old, new)@ pairs. + Rename [(T.Text, T.Text)] LogicalPlan + | -- | Skip the first N rows. + Offset Int LogicalPlan + | -- | A row-local eager op applied to each batch as it streams. + MapBatch T.Text ColumnEffect BatchFn LogicalPlan + | -- | An eager op applied once to the whole collected input. + Barrier T.Text ColumnEffect BatchFn LogicalPlan + | -- | An eager op over two whole collected inputs. + Barrier2 T.Text BatchFn2 LogicalPlan LogicalPlan | -- | Lift an already-loaded DataFrame into the lazy plan. SourceDF D.DataFrame deriving (Show) diff --git a/dataframe-lazy/src/DataFrame/Lazy/Internal/Optimizer.hs b/dataframe-lazy/src/DataFrame/Lazy/Internal/Optimizer.hs index 542a2321..8bdd3724 100644 --- a/dataframe-lazy/src/DataFrame/Lazy/Internal/Optimizer.hs +++ b/dataframe-lazy/src/DataFrame/Lazy/Internal/Optimizer.hs @@ -3,6 +3,7 @@ module DataFrame.Lazy.Internal.Optimizer (optimize) where +import qualified Data.List as L import qualified Data.Map as M import qualified Data.Set as S import qualified Data.Text as T @@ -33,7 +34,7 @@ fuseFilters (Aggregate keys aggs child) = Aggregate keys aggs (fuseFilters child) fuseFilters (Sort cols child) = Sort cols (fuseFilters child) fuseFilters (Limit n child) = Limit n (fuseFilters child) -fuseFilters leaf = leaf +fuseFilters plan = mapChildren fuseFilters plan -- | Logical AND of two @Bool@ expressions. andExpr :: E.Expr Bool -> E.Expr Bool -> E.Expr Bool @@ -60,6 +61,8 @@ pushPredicates (Filter p (Project cols child)) Project cols (pushPredicates (Filter p child)) | otherwise = Filter p (Project cols (pushPredicates child)) +pushPredicates (Filter p (MapBatch lbl eff f child)) + | commutes p eff = MapBatch lbl eff f (pushPredicates (Filter p child)) pushPredicates (Filter p child) = Filter p (pushPredicates child) pushPredicates (Project cols child) = Project cols (pushPredicates child) pushPredicates (Derive name expr child) = Derive name expr (pushPredicates child) @@ -69,7 +72,31 @@ pushPredicates (Aggregate keys aggs child) = Aggregate keys aggs (pushPredicates child) pushPredicates (Sort cols child) = Sort cols (pushPredicates child) pushPredicates (Limit n child) = Limit n (pushPredicates child) -pushPredicates leaf = leaf +pushPredicates plan = mapChildren pushPredicates plan + +-- | A filter moves below a row-local node it neither reads from nor feeds. +commutes :: E.Expr Bool -> ColumnEffect -> Bool +commutes p eff = case effectReads eff of + Nothing -> False + Just _ -> all (`notElem` effectWrites eff) (E.getColumns p) + +-- | Rebuild a node with an optimisation applied to each of its children. +mapChildren :: (LogicalPlan -> LogicalPlan) -> LogicalPlan -> LogicalPlan +mapChildren f plan = case plan of + Project cols c -> Project cols (f c) + Filter p c -> Filter p (f c) + Derive name e c -> Derive name e (f c) + Join jt ks l r -> Join jt ks (f l) (f r) + Aggregate ks aggs c -> Aggregate ks aggs (f c) + Sort ords c -> Sort ords (f c) + Limit n c -> Limit n (f c) + Rename pairs c -> Rename pairs (f c) + Offset n c -> Offset n (f c) + MapBatch lbl eff g c -> MapBatch lbl eff g (f c) + Barrier lbl eff g c -> Barrier lbl eff g (f c) + Barrier2 lbl g l r -> Barrier2 lbl g (f l) (f r) + Scan{} -> plan + SourceDF{} -> plan -- | Collect every column name referenced anywhere in the plan. referencedCols :: LogicalPlan -> Maybe (S.Set T.Text) @@ -92,8 +119,21 @@ referencedCols (Sort ords child) = (S.union (S.fromList (concatMap sortOrderColumns ords))) (referencedCols child) referencedCols (Limit _ child) = referencedCols child +referencedCols (Rename pairs child) = + fmap + (S.union (S.fromList (concatMap (\(o, n) -> [o, n]) pairs))) + (referencedCols child) +referencedCols (Offset _ child) = referencedCols child +referencedCols (MapBatch _ eff _ child) = effectCols eff child +referencedCols (Barrier _ eff _ child) = effectCols eff child +referencedCols (Barrier2{}) = Nothing referencedCols (SourceDF _) = Nothing +effectCols :: ColumnEffect -> LogicalPlan -> Maybe (S.Set T.Text) +effectCols eff child = do + rs <- effectReads eff + S.union (S.fromList (rs <> effectWrites eff)) <$> referencedCols child + someExprCols :: E.SomeExpr -> [T.Text] someExprCols (E.SomeExpr expr) = E.getColumns expr @@ -128,6 +168,15 @@ eliminateDeadColumns plan = go (referencedCols plan) plan ords (go (fmap (S.union (S.fromList (concatMap sortOrderColumns ords))) needed) child) go needed (Limit n child) = Limit n (go needed child) + go needed (Rename pairs child) = + let back c = maybe c fst (L.find ((== c) . snd) pairs) + in Rename pairs (go (fmap (S.map back) needed) child) + go needed (Offset n child) = Offset n (go needed child) + go needed (MapBatch lbl eff f child) = + MapBatch lbl eff f (go (throughEffect eff needed) child) + go needed (Barrier lbl eff f child) = + Barrier lbl eff f (go (throughEffect eff needed) child) + go _needed (Barrier2 lbl f l r) = Barrier2 lbl f (go Nothing l) (go Nothing r) go needed (Scan ds schema) = case needed of Nothing -> Scan ds schema @@ -135,6 +184,14 @@ eliminateDeadColumns plan = go (referencedCols plan) plan Scan ds (Schema (M.filterWithKey (\k _ -> k `S.member` cols) (elements schema))) go _ (SourceDF df) = SourceDF df +-- | Columns needed below a generic node, given those needed above it. +throughEffect :: ColumnEffect -> Maybe (S.Set T.Text) -> Maybe (S.Set T.Text) +throughEffect eff needed = do + rs <- effectReads eff + above <- needed + pure + ((above `S.difference` S.fromList (effectWrites eff)) `S.union` S.fromList rs) + -- | Lower the optimised logical plan to a physical plan. toPhysical :: Int -> LogicalPlan -> PhysicalPlan -- A Filter directly on a Scan is folded into the scan's pushdown predicate. @@ -184,4 +241,13 @@ toPhysical batchSz (Sort ords child) = PhysicalSort ords (toPhysical batchSz child) toPhysical batchSz (Limit n child) = PhysicalLimit n (toPhysical batchSz child) +toPhysical batchSz (Rename pairs child) = + PhysicalRename pairs (toPhysical batchSz child) +toPhysical batchSz (Offset n child) = PhysicalOffset n (toPhysical batchSz child) +toPhysical batchSz (MapBatch lbl _ f child) = + PhysicalMapBatch lbl f (toPhysical batchSz child) +toPhysical batchSz (Barrier lbl _ f child) = + PhysicalBarrier lbl f (toPhysical batchSz child) +toPhysical batchSz (Barrier2 lbl f l r) = + PhysicalBarrier2 lbl f (toPhysical batchSz l) (toPhysical batchSz r) toPhysical batchSz (SourceDF df) = PhysicalSourceDF batchSz df diff --git a/dataframe-lazy/src/DataFrame/Lazy/Internal/PhysicalPlan.hs b/dataframe-lazy/src/DataFrame/Lazy/Internal/PhysicalPlan.hs index 6fdb7f02..105430fe 100644 --- a/dataframe-lazy/src/DataFrame/Lazy/Internal/PhysicalPlan.hs +++ b/dataframe-lazy/src/DataFrame/Lazy/Internal/PhysicalPlan.hs @@ -3,7 +3,7 @@ module DataFrame.Lazy.Internal.PhysicalPlan where import qualified Data.Text as T import qualified DataFrame.Internal.DataFrame as D import qualified DataFrame.Internal.Expression as E -import DataFrame.Lazy.Internal.LogicalPlan (DataSource) +import DataFrame.Lazy.Internal.LogicalPlan (BatchFn, BatchFn2, DataSource) import DataFrame.Operations.Join (JoinType) import DataFrame.Operations.Permutation (SortOrder) import DataFrame.Schema (Schema) @@ -28,6 +28,11 @@ data PhysicalPlan | PhysicalHashAggregate [T.Text] [(T.Text, E.SomeExpr)] PhysicalPlan | PhysicalSort [SortOrder] PhysicalPlan | PhysicalLimit Int PhysicalPlan + | PhysicalRename [(T.Text, T.Text)] PhysicalPlan + | PhysicalOffset Int PhysicalPlan + | PhysicalMapBatch T.Text BatchFn PhysicalPlan + | PhysicalBarrier T.Text BatchFn PhysicalPlan + | PhysicalBarrier2 T.Text BatchFn2 PhysicalPlan PhysicalPlan | -- | Materialize child to a binary file on disk (used for build sides). PhysicalSpill PhysicalPlan FilePath | -- | Emit an already-loaded DataFrame as a stream of batches of size @n@. diff --git a/dataframe-lazy/src/DataFrame/Typed/Lazy.hs b/dataframe-lazy/src/DataFrame/Typed/Lazy.hs index 03d96ff6..9efae297 100644 --- a/dataframe-lazy/src/DataFrame/Typed/Lazy.hs +++ b/dataframe-lazy/src/DataFrame/Typed/Lazy.hs @@ -5,6 +5,8 @@ {-# LANGUAGE KindSignatures #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} -- | Typed lazy queries with the same operations as "DataFrame.Typed". module DataFrame.Typed.Lazy ( @@ -23,12 +25,70 @@ module DataFrame.Typed.Lazy ( -- * Schema-preserving operations filter, + filterBy, + filterJust, + filterNothing, + filterAllJust, + filterAllNothing, take, + takeLast, + drop, + dropLast, + range, + cube, sortBy, + topBy, + distinct, + sample, + shuffle, -- * Schema-modifying operations derive, + deriveMany, + deriveWithExpr, select, + exclude, + rename, + renameMany, + impute, + imputeWith, + insert, + insertColumn, + insertVector, + insertUnboxedVector, + insertWithDefault, + insertVectorWithDefault, + cloneColumn, + dropColumn, + replaceColumn, + apply, + applyInt, + applyDouble, + applyMany, + applyWhere, + applyAtIndex, + fold, + empty, + + -- * Merges and set operations + append, + (|||), + union, + intersect, + difference, + symmetricDifference, + + -- * Sampling and splitting + randomSplit, + kFolds, + selectRows, + stratifiedSample, + stratifiedSplit, + + -- * Frame-valued summaries (untyped results) + summarize, + describeColumns, + frequencies, -- * Aggregation groupBy, @@ -47,6 +107,54 @@ module DataFrame.Typed.Lazy ( -- * Execution run, + -- * Values (run the query) + safeApply, + null, + dimensions, + nRows, + nColumns, + columnNames, + valueCounts, + valueProportions, + mean, + meanMaybe, + median, + medianMaybe, + percentile, + genericPercentile, + standardDeviation, + skewness, + variance, + interQuartileRange, + sum, + correlation, + columnAsVector, + columnAsList, + columnAsIntVector, + columnAsDoubleVector, + columnAsFloatVector, + columnAsUnboxedVector, + toDoubleMatrix, + toFloatMatrix, + toIntMatrix, + toRecords, + + -- * Rendering and output + DisplayOptions (..), + defaultDisplayOptions, + display, + prettyPrint, + toCsv, + toCsv', + toMarkdown, + toMarkdown', + toSeparated, + writeCsv, + writeTsv, + writeSeparated, + writeParquet, + writeParquetWithOptions, + -- * Typed expressions (combinators live in "DataFrame.Typed.Functions") TExpr (..), TSortOrder (..), @@ -90,18 +198,65 @@ import Data.Kind (Type) import Data.Proxy (Proxy (..)) import qualified Data.Text as T import GHC.TypeLits (KnownSymbol, Symbol, symbolVal) -import Prelude hiding (filter, take) +import Prelude hiding (drop, filter, null, sum, take) +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU +import DataFrame.Display (DisplayOptions (..), defaultDisplayOptions) +import DataFrame.Errors (DataFrameException) import DataFrame.Expression.Operators ((|>)) +import qualified DataFrame.Functions as DF import DataFrame.IO.CSV (CsvReader) +import DataFrame.IO.Parquet (ParquetWriteOptions) import qualified DataFrame.Internal.Column as C +import qualified DataFrame.Internal.Expression as E +import DataFrame.Internal.Expression.Operators.Nullable (BaseType) +import qualified DataFrame.Lazy.Internal.Collect as LC import DataFrame.Lazy.Internal.DataFrame (LazyDataFrame) import qualified DataFrame.Lazy.Internal.DataFrame as L import DataFrame.Operations.Join (JoinType (..)) +import DataFrame.Operations.Subset (SplittableGen) +import DataFrame.Operations.Transformations (ImputeOp) import DataFrame.Schema (RuntimeSchema (..)) +import qualified DataFrame.Typed.Access as DTA import DataFrame.Typed.Aggregate (as) -import DataFrame.Typed.Expr -import DataFrame.Typed.Freeze (unsafeFreeze) +import qualified DataFrame.Typed.Apply as DTApply +import DataFrame.Typed.Expr ( + caseWhen, + col, + ifThenElse, + lit, + orElse, + (.&&), + (.&&.), + (.*), + (.*.), + (.+), + (.+.), + (.-), + (.-.), + (./), + (./.), + (./=), + (./=.), + (.<), + (.<.), + (.<=), + (.<=.), + (.==), + (.==.), + (.>), + (.>.), + (.>=), + (.>=.), + (.^), + (.^.), + (.^^), + (.^^.), + (.||), + (.||.), + ) +import DataFrame.Typed.Freeze (thaw, unsafeFreeze) import DataFrame.Typed.Join ( Fsts, JoinKeys, @@ -111,8 +266,11 @@ import DataFrame.Typed.Join ( SwapPairs, ) import DataFrame.Typed.Operations (AllKnownPairs (..), toUntypedSort) +import qualified DataFrame.Typed.Operations as DT +import DataFrame.Typed.Record (HasSchema (..)) import DataFrame.Typed.Schema import DataFrame.Typed.Types +import System.Random (RandomGen) -- | A lazy query with compile-time schema tracking. newtype TypedLazyDataFrame (cols :: [(Symbol, Type)]) = TLD {_unTLD :: LazyDataFrame} @@ -265,3 +423,601 @@ fullOuterJoin = join @'FULL_OUTER @keys -- | Execute the lazy query and return a typed DataFrame. run :: TypedLazyDataFrame cols -> IO (TypedDataFrame cols) run (TLD ldf) = unsafeFreeze <$> L.run ldf + +-- * Helpers + +nameOf :: forall (n :: Symbol). (KnownSymbol n) => T.Text +nameOf = T.pack (symbolVal (Proxy @n)) + +-- | Run the query and apply a typed eager function to the result. +collectWith :: (TypedDataFrame cols -> r) -> TypedLazyDataFrame cols -> IO r +collectWith f tld = f <$> run tld + +-- * Row operations + +filterBy :: + (C.Columnable a) => + (a -> Bool) -> + TExpr cols a -> + TypedLazyDataFrame cols -> + TypedLazyDataFrame cols +filterBy p (TExpr e) (TLD ldf) = TLD (L.filterBy p e ldf) + +filterJust :: + forall name cols. + (KnownSymbol name, AssertPresent name cols) => + TypedLazyDataFrame cols -> TypedLazyDataFrame (StripMaybeAt name cols) +filterJust (TLD ldf) = TLD (L.filterJust (nameOf @name) ldf) + +filterNothing :: + forall name cols. + (KnownSymbol name, AssertPresent name cols) => + TypedLazyDataFrame cols -> TypedLazyDataFrame cols +filterNothing (TLD ldf) = TLD (L.filterNothing (nameOf @name) ldf) + +filterAllJust :: + TypedLazyDataFrame cols -> TypedLazyDataFrame (StripAllMaybe cols) +filterAllJust (TLD ldf) = TLD (L.filterAllJust ldf) + +filterAllNothing :: TypedLazyDataFrame cols -> TypedLazyDataFrame cols +filterAllNothing (TLD ldf) = TLD (L.filterAllNothing ldf) + +takeLast :: Int -> TypedLazyDataFrame cols -> TypedLazyDataFrame cols +takeLast n (TLD ldf) = TLD (L.takeLast n ldf) + +drop :: Int -> TypedLazyDataFrame cols -> TypedLazyDataFrame cols +drop n (TLD ldf) = TLD (L.drop n ldf) + +dropLast :: Int -> TypedLazyDataFrame cols -> TypedLazyDataFrame cols +dropLast n (TLD ldf) = TLD (L.dropLast n ldf) + +range :: (Int, Int) -> TypedLazyDataFrame cols -> TypedLazyDataFrame cols +range r (TLD ldf) = TLD (L.range r ldf) + +cube :: (Int, Int) -> TypedLazyDataFrame cols -> TypedLazyDataFrame cols +cube c (TLD ldf) = TLD (L.cube c ldf) + +topBy :: + Int -> + Int -> + [TSortOrder cols] -> + TypedLazyDataFrame cols -> + TypedLazyDataFrame cols +topBy offset count ords (TLD ldf) = + TLD (L.topBy offset count (map toUntypedSort ords) ldf) + +distinct :: TypedLazyDataFrame cols -> TypedLazyDataFrame cols +distinct (TLD ldf) = TLD (L.distinct ldf) + +sample :: + (RandomGen g) => + g -> Double -> TypedLazyDataFrame cols -> TypedLazyDataFrame cols +sample g p (TLD ldf) = TLD (L.sample g p ldf) + +shuffle :: + (RandomGen g) => g -> TypedLazyDataFrame cols -> TypedLazyDataFrame cols +shuffle g (TLD ldf) = TLD (L.shuffle g ldf) + +-- * Column operations + +deriveMany :: + forall new cols. + (AssertAllAbsent (Reverse new) cols) => + (TAgg '[] cols '[] -> TAgg '[] cols new) -> + TypedLazyDataFrame cols -> + TypedLazyDataFrame (Append cols (Reverse new)) +deriveMany build (TLD ldf) = + TLD (L.deriveMany (taggToNamedExprs (build TAggNil)) ldf) + +deriveWithExpr :: + forall name a cols. + (KnownSymbol name, C.Columnable a, AssertAbsent name cols) => + TExpr cols a -> + TypedLazyDataFrame cols -> + ( TExpr (Snoc cols '(name, a)) a + , TypedLazyDataFrame (Snoc cols '(name, a)) + ) +deriveWithExpr e tld = (TExpr (E.Col (nameOf @name)), derive @name e tld) + +exclude :: + forall (names :: [Symbol]) cols. + (AllKnownSymbol names) => + TypedLazyDataFrame cols -> TypedLazyDataFrame (ExcludeSchema names cols) +exclude (TLD ldf) = TLD (L.exclude (symbolVals @names) ldf) + +rename :: + forall old new cols. + (KnownSymbol old, KnownSymbol new) => + TypedLazyDataFrame cols -> TypedLazyDataFrame (RenameInSchema old new cols) +rename (TLD ldf) = TLD (L.rename (nameOf @old) (nameOf @new) ldf) + +renameMany :: + forall (pairs :: [(Symbol, Symbol)]) cols. + (AllKnownPairs pairs) => + TypedLazyDataFrame cols -> TypedLazyDataFrame (RenameManyInSchema pairs cols) +renameMany (TLD ldf) = TLD (L.renameMany (pairVals @pairs) ldf) + +impute :: + forall name a cols. + (KnownSymbol name, C.Columnable a, Maybe a ~ Lookup name cols) => + a -> TypedLazyDataFrame cols -> TypedLazyDataFrame (Impute name cols) +impute v (TLD ldf) = + TLD + (L.derive (nameOf @name) (DF.fromMaybe v (DF.col @(Maybe a) (nameOf @name))) ldf) + +imputeWith :: + (ImputeOp a, C.Columnable (BaseType a)) => + (TExpr cols (BaseType a) -> TExpr cols (BaseType a)) -> + TExpr cols a -> + TypedLazyDataFrame cols -> + TypedLazyDataFrame cols +imputeWith f (TExpr e) (TLD ldf) = TLD (L.imputeWith (unTExpr . f . TExpr) e ldf) + +insert :: + forall name a cols t. + (KnownSymbol name, C.Columnable a, Foldable t, AssertAbsent name cols) => + t a -> TypedLazyDataFrame cols -> TypedLazyDataFrame ('(name, a) ': cols) +insert xs (TLD ldf) = TLD (L.insert (nameOf @name) xs ldf) + +insertColumn :: + forall name a cols. + (KnownSymbol name, C.Columnable a, AssertAbsent name cols) => + C.Column -> TypedLazyDataFrame cols -> TypedLazyDataFrame ('(name, a) ': cols) +insertColumn c (TLD ldf) = TLD (L.insertColumn (nameOf @name) c ldf) + +insertVector :: + forall name a cols. + (KnownSymbol name, C.Columnable a, AssertAbsent name cols) => + V.Vector a -> TypedLazyDataFrame cols -> TypedLazyDataFrame ('(name, a) ': cols) +insertVector xs (TLD ldf) = TLD (L.insertVector (nameOf @name) xs ldf) + +insertUnboxedVector :: + forall name a cols. + (KnownSymbol name, C.Columnable a, VU.Unbox a, AssertAbsent name cols) => + VU.Vector a -> + TypedLazyDataFrame cols -> + TypedLazyDataFrame ('(name, a) ': cols) +insertUnboxedVector xs (TLD ldf) = TLD (L.insertUnboxedVector (nameOf @name) xs ldf) + +insertWithDefault :: + forall name a cols t. + (KnownSymbol name, C.Columnable a, Foldable t, AssertAbsent name cols) => + a -> t a -> TypedLazyDataFrame cols -> TypedLazyDataFrame ('(name, a) ': cols) +insertWithDefault d xs (TLD ldf) = TLD (L.insertWithDefault d (nameOf @name) xs ldf) + +insertVectorWithDefault :: + forall name a cols. + (KnownSymbol name, C.Columnable a, AssertAbsent name cols) => + a -> + V.Vector a -> + TypedLazyDataFrame cols -> + TypedLazyDataFrame ('(name, a) ': cols) +insertVectorWithDefault d xs (TLD ldf) = + TLD (L.insertVectorWithDefault d (nameOf @name) xs ldf) + +cloneColumn :: + forall old new cols. + ( KnownSymbol old + , KnownSymbol new + , AssertPresent old cols + , AssertAbsent new cols + ) => + TypedLazyDataFrame cols -> TypedLazyDataFrame ('(new, Lookup old cols) ': cols) +cloneColumn (TLD ldf) = TLD (L.cloneColumn (nameOf @old) (nameOf @new) ldf) + +dropColumn :: + forall name cols. + (KnownSymbol name, AssertPresent name cols) => + TypedLazyDataFrame cols -> TypedLazyDataFrame (RemoveColumn name cols) +dropColumn (TLD ldf) = TLD (L.exclude [nameOf @name] ldf) + +replaceColumn :: + forall name a cols. + ( KnownSymbol name + , C.Columnable a + , a ~ SafeLookup name cols + , AssertPresent name cols + ) => + TExpr cols a -> TypedLazyDataFrame cols -> TypedLazyDataFrame cols +replaceColumn (TExpr e) (TLD ldf) = TLD (L.derive (nameOf @name) e ldf) + +apply :: + forall name a b cols. + ( KnownSymbol name + , a ~ SafeLookup name cols + , C.Columnable a + , C.Columnable b + , AssertPresent name cols + ) => + (a -> b) -> + TypedLazyDataFrame cols -> + TypedLazyDataFrame (SetColumnType name b cols) +apply f (TLD ldf) = TLD (L.apply f (nameOf @name) ldf) + +applyInt :: + forall name b cols. + ( KnownSymbol name + , Int ~ SafeLookup name cols + , C.Columnable b + , AssertPresent name cols + ) => + (Int -> b) -> + TypedLazyDataFrame cols -> + TypedLazyDataFrame (SetColumnType name b cols) +applyInt = apply @name + +applyDouble :: + forall name b cols. + ( KnownSymbol name + , Double ~ SafeLookup name cols + , C.Columnable b + , AssertPresent name cols + ) => + (Double -> b) -> + TypedLazyDataFrame cols -> + TypedLazyDataFrame (SetColumnType name b cols) +applyDouble = apply @name + +applyMany :: + forall (names :: [Symbol]) a cols. + (AllKnownSymbol names, C.Columnable a, AssertAllColumnsHaveType names a cols) => + (a -> a) -> TypedLazyDataFrame cols -> TypedLazyDataFrame cols +applyMany f (TLD ldf) = TLD (L.applyMany f (symbolVals @names) ldf) + +applyWhere :: + forall filterName targetName a b cols. + ( KnownSymbol filterName + , KnownSymbol targetName + , a ~ SafeLookup filterName cols + , b ~ SafeLookup targetName cols + , C.Columnable a + , C.Columnable b + , AssertPresent filterName cols + , AssertPresent targetName cols + ) => + (a -> Bool) -> (b -> b) -> TypedLazyDataFrame cols -> TypedLazyDataFrame cols +applyWhere p f (TLD ldf) = + TLD (L.applyWhere p (nameOf @filterName) f (nameOf @targetName) ldf) + +applyAtIndex :: + forall name a cols. + ( KnownSymbol name + , a ~ SafeLookup name cols + , C.Columnable a + , AssertPresent name cols + ) => + Int -> (a -> a) -> TypedLazyDataFrame cols -> TypedLazyDataFrame cols +applyAtIndex i f (TLD ldf) = TLD (L.applyAtIndex i f (nameOf @name) ldf) + +fold :: + (a -> TypedLazyDataFrame cols -> TypedLazyDataFrame cols) -> + [a] -> + TypedLazyDataFrame cols -> + TypedLazyDataFrame cols +fold f xs tld = foldl (flip f) tld xs + +-- | The query with no rows and no columns. +empty :: TypedLazyDataFrame '[] +empty = TLD (L.fromDataFrame (thaw DT.empty)) + +-- * Merges and set operations + +append :: + TypedLazyDataFrame cols -> TypedLazyDataFrame cols -> TypedLazyDataFrame cols +append (TLD a) (TLD b) = TLD (a <> b) + +(|||) :: + (AssertDisjoint left right) => + TypedLazyDataFrame left -> + TypedLazyDataFrame right -> + TypedLazyDataFrame (Append left right) +TLD a ||| TLD b = TLD (a L.||| b) + +union :: + TypedLazyDataFrame cols -> TypedLazyDataFrame cols -> TypedLazyDataFrame cols +union (TLD a) (TLD b) = TLD (L.union a b) + +intersect :: + TypedLazyDataFrame cols -> TypedLazyDataFrame cols -> TypedLazyDataFrame cols +intersect (TLD a) (TLD b) = TLD (L.intersect a b) + +difference :: + TypedLazyDataFrame cols -> TypedLazyDataFrame cols -> TypedLazyDataFrame cols +difference (TLD a) (TLD b) = TLD (L.difference a b) + +symmetricDifference :: + TypedLazyDataFrame cols -> TypedLazyDataFrame cols -> TypedLazyDataFrame cols +symmetricDifference (TLD a) (TLD b) = TLD (L.symmetricDifference a b) + +-- * Sampling and splitting + +randomSplit :: + (RandomGen g) => + g -> + Double -> + TypedLazyDataFrame cols -> + (TypedLazyDataFrame cols, TypedLazyDataFrame cols) +randomSplit g p (TLD ldf) = let (a, b) = L.randomSplit g p ldf in (TLD a, TLD b) + +kFolds :: + (RandomGen g) => + g -> Int -> TypedLazyDataFrame cols -> [TypedLazyDataFrame cols] +kFolds g k (TLD ldf) = map TLD (L.kFolds g k ldf) + +selectRows :: [Int] -> TypedLazyDataFrame cols -> TypedLazyDataFrame cols +selectRows ixs (TLD ldf) = TLD (L.selectRows ixs ldf) + +stratifiedSample :: + (SplittableGen g, C.Columnable a) => + g -> + Double -> + TExpr cols a -> + TypedLazyDataFrame cols -> + TypedLazyDataFrame cols +stratifiedSample g p (TExpr e) (TLD ldf) = TLD (L.stratifiedSample g p e ldf) + +stratifiedSplit :: + (SplittableGen g, C.Columnable a) => + g -> + Double -> + TExpr cols a -> + TypedLazyDataFrame cols -> + (TypedLazyDataFrame cols, TypedLazyDataFrame cols) +stratifiedSplit g p (TExpr e) (TLD ldf) = + let (a, b) = L.stratifiedSplit g p e ldf in (TLD a, TLD b) + +-- * Frame-valued summaries + +summarize :: TypedLazyDataFrame cols -> LazyDataFrame +summarize (TLD ldf) = L.summarize ldf + +describeColumns :: TypedLazyDataFrame cols -> LazyDataFrame +describeColumns (TLD ldf) = L.describeColumns ldf + +frequencies :: + (C.Columnable a, Ord a) => + TExpr cols a -> TypedLazyDataFrame cols -> LazyDataFrame +frequencies (TExpr e) (TLD ldf) = L.frequencies e ldf + +-- * Values + +safeApply :: + forall name a b cols. + ( KnownSymbol name + , a ~ SafeLookup name cols + , C.Columnable a + , C.Columnable b + , AssertPresent name cols + ) => + (a -> b) -> + TypedLazyDataFrame cols -> + IO (Either DataFrameException (TypedDataFrame (SetColumnType name b cols))) +safeApply f = collectWith (DTApply.safeApply @name f) + +null :: TypedLazyDataFrame cols -> IO Bool +null (TLD ldf) = LC.null ldf + +dimensions :: TypedLazyDataFrame cols -> IO (Int, Int) +dimensions (TLD ldf) = LC.dimensions ldf + +nRows :: TypedLazyDataFrame cols -> IO Int +nRows (TLD ldf) = LC.nRows ldf + +nColumns :: TypedLazyDataFrame cols -> IO Int +nColumns (TLD ldf) = LC.nColumns ldf + +columnNames :: TypedLazyDataFrame cols -> IO [T.Text] +columnNames (TLD ldf) = LC.columnNames ldf + +valueCounts :: + (Ord a, C.Columnable a) => + TExpr cols a -> TypedLazyDataFrame cols -> IO [(a, Int)] +valueCounts (TExpr e) (TLD ldf) = LC.valueCounts e ldf + +valueProportions :: + (Ord a, C.Columnable a) => + TExpr cols a -> TypedLazyDataFrame cols -> IO [(a, Double)] +valueProportions (TExpr e) (TLD ldf) = LC.valueProportions e ldf + +mean :: + (C.Columnable a, Real a, VU.Unbox a) => + TExpr cols a -> TypedLazyDataFrame cols -> IO Double +mean (TExpr e) (TLD ldf) = LC.mean e ldf + +meanMaybe :: + (C.Columnable a, Real a) => + TExpr cols (Maybe a) -> TypedLazyDataFrame cols -> IO Double +meanMaybe (TExpr e) (TLD ldf) = LC.meanMaybe e ldf + +median :: + (C.Columnable a, Real a, VU.Unbox a) => + TExpr cols a -> TypedLazyDataFrame cols -> IO Double +median (TExpr e) (TLD ldf) = LC.median e ldf + +medianMaybe :: + (C.Columnable a, Real a) => + TExpr cols (Maybe a) -> TypedLazyDataFrame cols -> IO Double +medianMaybe (TExpr e) (TLD ldf) = LC.medianMaybe e ldf + +percentile :: + (C.Columnable a, Real a, VU.Unbox a) => + Int -> TExpr cols a -> TypedLazyDataFrame cols -> IO Double +percentile p (TExpr e) (TLD ldf) = LC.percentile p e ldf + +genericPercentile :: + (C.Columnable a, Ord a) => + Int -> TExpr cols a -> TypedLazyDataFrame cols -> IO a +genericPercentile p (TExpr e) (TLD ldf) = LC.genericPercentile p e ldf + +standardDeviation :: + (C.Columnable a, Real a, VU.Unbox a) => + TExpr cols a -> TypedLazyDataFrame cols -> IO Double +standardDeviation (TExpr e) (TLD ldf) = LC.standardDeviation e ldf + +skewness :: + (C.Columnable a, Real a, VU.Unbox a) => + TExpr cols a -> TypedLazyDataFrame cols -> IO Double +skewness (TExpr e) (TLD ldf) = LC.skewness e ldf + +variance :: + (C.Columnable a, Real a, VU.Unbox a) => + TExpr cols a -> TypedLazyDataFrame cols -> IO Double +variance (TExpr e) (TLD ldf) = LC.variance e ldf + +interQuartileRange :: + (C.Columnable a, Real a, VU.Unbox a) => + TExpr cols a -> TypedLazyDataFrame cols -> IO Double +interQuartileRange (TExpr e) (TLD ldf) = LC.interQuartileRange e ldf + +sum :: + (C.Columnable a, Num a) => TExpr cols a -> TypedLazyDataFrame cols -> IO a +sum (TExpr e) (TLD ldf) = LC.sum e ldf + +correlation :: + forall c1 c2 a b cols. + ( KnownSymbol c1 + , KnownSymbol c2 + , a ~ SafeLookup c1 cols + , b ~ SafeLookup c2 cols + , C.Columnable a + , C.Columnable b + , Real a + , Real b + , VU.Unbox a + , VU.Unbox b + , AssertPresent c1 cols + , AssertPresent c2 cols + ) => + TypedLazyDataFrame cols -> IO (Maybe Double) +correlation (TLD ldf) = LC.correlation (nameOf @c1) (nameOf @c2) ldf + +columnAsVector :: + forall name cols a. + ( KnownSymbol name + , a ~ SafeLookup name cols + , C.Columnable a + , AssertPresent name cols + ) => + TypedLazyDataFrame cols -> IO (V.Vector a) +columnAsVector = collectWith (DTA.columnAsVector @name) + +columnAsList :: + forall name cols a. + ( KnownSymbol name + , a ~ SafeLookup name cols + , C.Columnable a + , AssertPresent name cols + ) => + TypedLazyDataFrame cols -> IO [a] +columnAsList = collectWith (DTA.columnAsList @name) + +columnAsIntVector :: + forall name cols a. + ( KnownSymbol name + , a ~ SafeLookup name cols + , C.Columnable a + , AssertRealColumn "columnAsIntVector" name a + , Real a + , VU.Unbox a + , AssertPresent name cols + ) => + TypedLazyDataFrame cols -> IO (VU.Vector Int) +columnAsIntVector = collectWith (DTA.columnAsIntVector @name) + +columnAsDoubleVector :: + forall name cols a. + ( KnownSymbol name + , a ~ SafeLookup name cols + , C.Columnable a + , AssertRealColumn "columnAsDoubleVector" name a + , Real a + , VU.Unbox a + , AssertPresent name cols + ) => + TypedLazyDataFrame cols -> IO (VU.Vector Double) +columnAsDoubleVector = collectWith (DTA.columnAsDoubleVector @name) + +columnAsFloatVector :: + forall name cols a. + ( KnownSymbol name + , a ~ SafeLookup name cols + , C.Columnable a + , AssertRealColumn "columnAsFloatVector" name a + , Real a + , VU.Unbox a + , AssertPresent name cols + ) => + TypedLazyDataFrame cols -> IO (VU.Vector Float) +columnAsFloatVector = collectWith (DTA.columnAsFloatVector @name) + +columnAsUnboxedVector :: + forall name cols a. + ( KnownSymbol name + , a ~ SafeLookup name cols + , C.Columnable a + , VU.Unbox a + , AssertPresent name cols + ) => + TypedLazyDataFrame cols -> IO (VU.Vector a) +columnAsUnboxedVector = collectWith (DTA.columnAsUnboxedVector @name) + +toDoubleMatrix :: + (AllColumnsReal "toDoubleMatrix" cols) => + TypedLazyDataFrame cols -> IO (V.Vector (VU.Vector Double)) +toDoubleMatrix = collectWith DTA.toDoubleMatrix + +toFloatMatrix :: + (AllColumnsReal "toFloatMatrix" cols) => + TypedLazyDataFrame cols -> IO (V.Vector (VU.Vector Float)) +toFloatMatrix = collectWith DTA.toFloatMatrix + +toIntMatrix :: + (AllColumnsReal "toIntMatrix" cols) => + TypedLazyDataFrame cols -> IO (V.Vector (VU.Vector Int)) +toIntMatrix = collectWith DTA.toIntMatrix + +-- | Parse records out of the result. +toRecords :: + forall a. + (HasSchema a) => TypedLazyDataFrame (Schema a) -> IO (Either T.Text [a]) +toRecords (TLD ldf) = LC.toRecords ldf + +-- * Rendering and output + +display :: DisplayOptions -> TypedLazyDataFrame cols -> IO () +display opts (TLD ldf) = LC.display opts ldf + +-- | Render a typed expression as source-like text. +prettyPrint :: TExpr cols a -> String +prettyPrint (TExpr e) = E.prettyPrint e + +toCsv :: TypedLazyDataFrame cols -> IO T.Text +toCsv (TLD ldf) = LC.toCsv ldf + +toCsv' :: TypedLazyDataFrame cols -> IO String +toCsv' (TLD ldf) = LC.toCsv' ldf + +toMarkdown :: TypedLazyDataFrame cols -> IO T.Text +toMarkdown (TLD ldf) = LC.toMarkdown ldf + +toMarkdown' :: TypedLazyDataFrame cols -> IO String +toMarkdown' (TLD ldf) = LC.toMarkdown' ldf + +toSeparated :: Char -> TypedLazyDataFrame cols -> IO T.Text +toSeparated sep (TLD ldf) = LC.toSeparated sep ldf + +writeCsv :: FilePath -> TypedLazyDataFrame cols -> IO () +writeCsv path (TLD ldf) = LC.writeCsv path ldf + +writeTsv :: FilePath -> TypedLazyDataFrame cols -> IO () +writeTsv path (TLD ldf) = LC.writeTsv path ldf + +writeSeparated :: Char -> FilePath -> TypedLazyDataFrame cols -> IO () +writeSeparated sep path (TLD ldf) = LC.writeSeparated sep path ldf + +writeParquet :: FilePath -> TypedLazyDataFrame cols -> IO () +writeParquet path (TLD ldf) = LC.writeParquet path ldf + +writeParquetWithOptions :: + ParquetWriteOptions -> FilePath -> TypedLazyDataFrame cols -> IO () +writeParquetWithOptions opts path (TLD ldf) = LC.writeParquetWithOptions opts path ldf diff --git a/docs/base_scripts/base_readme.md b/docs/base_scripts/base_readme.md index 87bb93d2..47cd6eb1 100644 --- a/docs/base_scripts/base_readme.md +++ b/docs/base_scripts/base_readme.md @@ -401,6 +401,23 @@ D.take 10 lazyResult |> D.toMarkdown' The optimizer pushes the filter into the scan, drops unreferenced columns before reading, and stops pulling batches once 1000 rows have been collected. +### Switching between APIs + +The eager, typed, lazy and typed-lazy APIs share operation names and argument +order, so moving a pipeline between them is mostly a change of imports: + +| | Frame ops (`D`) | Expressions (`F`) | +|---|---|---| +| Eager | `DataFrame` | `DataFrame.Functions` | +| Typed | `DataFrame.Typed` | `DataFrame.Typed.Functions` | +| Lazy | `DataFrame.Lazy` | `DataFrame.Functions` | +| Typed lazy | `DataFrame.Typed.Lazy` | `DataFrame.Typed.Functions` | + +A lazy pipeline starts from `scanCsv`/`scanParquet`/`fromDataFrame` and ends with +`run`; ops that return a value (`mean`, `nRows`, `writeCsv`, …) run the query +and return `IO`. Typed APIs name columns at the type level (`col @"x"`, +`select @'["x"]`) where untyped ones take text. + ## Documentation * User guide: https://dataframe.readthedocs.io/en/latest/ diff --git a/tests/LazyParity.hs b/tests/LazyParity.hs index d9976679..2b2a4c2c 100644 --- a/tests/LazyParity.hs +++ b/tests/LazyParity.hs @@ -8,7 +8,7 @@ import qualified Data.Map.Strict as M import Data.Text (Text) import qualified Data.Text as T import qualified DataFrame as D -import DataFrame.Expression.Operators (as, (|>)) +import DataFrame.Expression.Operators (as, (.=), (.>.), (|>)) import qualified DataFrame.Functions as F import qualified DataFrame.IO.CSV as Csv import qualified DataFrame.Internal.Column as DI @@ -211,9 +211,76 @@ streamingJoinParity = (L.fromDataFrame renamedCustomers) ) +-- | New plan nodes under a scan, where pruning and pushdown are active. +optimizedNodesParity :: Test +optimizedNodesParity = + TestCase $ + withCsv (ordersFrame 3_000) $ \ordersPath -> do + df <- Csv.readCsvWithSchema ordersSchema ordersPath + let scan = (L.scanCsv ordersSchema (T.pack ordersPath)){L.batchSize = 256} + amount = F.col @Double "amount" + check lbl eager lazy = do + actual <- L.run lazy + assertEqual lbl (show eager) (show actual) + check + "rename, filter on new name, select" + ( D.select ["amt", "order_id"] $ + D.filter (F.col @Double "amt" .>. 1000) $ + D.rename "amount" "amt" df + ) + ( L.select ["amt", "order_id"] $ + L.filter (F.col @Double "amt" .>. 1000) $ + L.rename "amount" "amt" scan + ) + check + "filter above drop stays above" + (D.filter (amount .>. 100) (D.drop 50 df)) + (L.filter (amount .>. 100) (L.drop 50 scan)) + check + "filter above apply moves below it" + ( D.select ["order_id", "discount"] $ + D.filter (amount .>. 4000) $ + D.apply (* (2 :: Double)) "discount" df + ) + ( L.select ["order_id", "discount"] $ + L.filter (amount .>. 4000) $ + L.apply (* (2 :: Double)) "discount" scan + ) + check + "filter on applied column stays above" + ( D.filter + (F.col @Double "discount" .>. 2) + (D.apply (* (2 :: Double)) "discount" df) + ) + ( L.filter + (F.col @Double "discount" .>. 2) + (L.apply (* (2 :: Double)) "discount" scan) + ) + check + "exclude then select" + (D.select ["order_id"] (D.exclude ["discount"] df)) + (L.select ["order_id"] (L.exclude ["discount"] scan)) + check + "filter above takeLast stays above" + (D.filter (amount .>. 10) (D.takeLast 100 df)) + (L.filter (amount .>. 10) (L.takeLast 100 scan)) + check + "deriveMany then select a derived column" + ( D.select ["x"] $ + D.deriveMany ["x" .= (amount * 2), "y" .= amount] df + ) + ( L.select ["x"] $ + L.deriveMany ["x" .= (amount * 2), "y" .= amount] scan + ) + check + "append two scans" + (df <> D.take 10 df) + (scan <> L.take 10 scan) + tests :: [Test] tests = - [ joinPipelineParity + [ optimizedNodesParity + , joinPipelineParity , groupByPipelineParity , streamingMeanParity , streamingJoinParity diff --git a/tests/LazyParquet.hs b/tests/LazyParquet.hs index 210dd639..b45fcd92 100644 --- a/tests/LazyParquet.hs +++ b/tests/LazyParquet.hs @@ -119,9 +119,20 @@ sortAndLimit = assertEqual "sortAndLimit rows" 3 (fst (D.dimensions actual)) ) +-- | A scan returns its columns in file order, like an eager read. +fileColumnOrder :: Test +fileColumnOrder = + TestCase + ( do + actual <- L.run (L.scanParquet allTypesSchema (T.pack plainPath)) + eager <- D.readParquet plainPath + assertEqual "column order" (D.columnNames eager) (D.columnNames actual) + ) + tests :: [Test] tests = - [ basicScan + [ fileColumnOrder + , basicScan , columnProjection , filterPushdown , filterAndProject diff --git a/tests/LazyProjection.hs b/tests/LazyProjection.hs index d62fe544..e4858c93 100644 --- a/tests/LazyProjection.hs +++ b/tests/LazyProjection.hs @@ -40,12 +40,12 @@ twoCols :: Schema twoCols = Schema $ M.fromList [("id", schemaType @Int), ("dob", schemaType @Int)] --- | The scan keeps only the schema's columns, whatever their header position. +-- | The scan keeps only the schema's columns, in header order like eager. scanProjects :: Test scanProjects = TestLabel "scan_projects" $ TestCase $ withCsv ',' $ \csvPath -> do df <- L.run (L.scanCsv twoCols (T.pack csvPath)) assertEqual "dims" (2, 2) (D.dimensions df) - assertEqual "columns" ["dob", "id"] (D.columnNames df) + assertEqual "columns" ["id", "dob"] (D.columnNames df) -- | Same schema, same projection, through the SIMD reader. scanProjectsWithFastReader :: Test @@ -55,7 +55,7 @@ scanProjectsWithFastReader = L.run (L.scanCsvWith Fast.fastReadCsvWithOpts twoCols (T.pack csvPath)) assertEqual "dims" (2, 2) (D.dimensions df) - assertEqual "columns" ["dob", "id"] (D.columnNames df) + assertEqual "columns" ["id", "dob"] (D.columnNames df) -- | Both readers must agree on what a scan produces. readersAgree :: Test @@ -72,7 +72,7 @@ scanHonoursSeparator = -- With the separator ignored, the schema's columns are not found. df <- L.run (L.scanSeparated ';' twoCols (T.pack csvPath)) assertEqual "dims" (2, 2) (D.dimensions df) - assertEqual "columns" ["dob", "id"] (D.columnNames df) + assertEqual "columns" ["id", "dob"] (D.columnNames df) -- | A schema naming a column the file lacks is an error, not a short frame. scanMissingColumn :: Test diff --git a/tests/Parity/FourWay.hs b/tests/Parity/FourWay.hs index 2234c605..8452ceac 100644 --- a/tests/Parity/FourWay.hs +++ b/tests/Parity/FourWay.hs @@ -6,6 +6,7 @@ module Parity.FourWay (tests) where import qualified Data.Text as T +import System.Random (StdGen, mkStdGen) import qualified DataFrame as D import qualified DataFrame.Functions as F @@ -54,6 +55,22 @@ usersDF = , ("name", DI.fromList ["ann", "bo", "cy" :: T.Text]) ] +type Nullable = '[ '("k", Int), '("v", Maybe Double), '("w", Maybe Int)] + +nullableDF :: D.DataFrame +nullableDF = + D.fromColumns + [ ("k", DI.fromList [1, 2, 3, 4, 5 :: Int]) + , ("v", DI.fromList [Just 1.5, Nothing, Just 3, Nothing, Just 5 :: Maybe Double]) + , ("w", DI.fromList [Nothing, Nothing, Just 3, Nothing, Just 1 :: Maybe Int]) + ] + +nullableT :: DT.TypedDataFrame Nullable +nullableT = either (error . show) id (DT.freezeWithError nullableDF) + +tln :: TL.TypedLazyDataFrame Nullable +tln = TL.fromDataFrame nullableT + ordersT :: DT.TypedDataFrame Orders ordersT = either (error . show) id (DT.freezeWithError ordersDF) @@ -72,16 +89,29 @@ fourWay :: TL.TypedLazyDataFrame cols -> Test fourWay lbl eager typed lazy typedLazy = TestCase $ do - assertEqual (lbl ++ ": typed == eager") eager (DT.thaw typed) + sameFrame (lbl ++ ": typed == eager") eager (DT.thaw typed) lazyDf <- L.run lazy - assertEqual (lbl ++ ": lazy == eager") eager lazyDf + sameFrame (lbl ++ ": lazy == eager") eager lazyDf typedLazyDf <- TL.run typedLazy - assertEqual (lbl ++ ": typed lazy == eager") eager (DT.thaw typedLazyDf) + sameFrame (lbl ++ ": typed lazy == eager") eager (DT.thaw typedLazyDf) + +-- | Equal values and equal column order ('==' ignores the order). +sameFrame :: String -> D.DataFrame -> D.DataFrame -> Assertion +sameFrame lbl expected actual = + assertEqual + lbl + (D.columnNames expected, expected) + (D.columnNames actual, actual) + +-- | Two-row batches, so streaming ops cross batch boundaries. +lazyOf :: D.DataFrame -> L.LazyDataFrame +lazyOf df = (L.fromDataFrame df){L.batchSize = 2} -lo, lc, lu :: L.LazyDataFrame -lo = L.fromDataFrame ordersDF -lc = L.fromDataFrame customersDF -lu = L.fromDataFrame usersDF +lo, lc, lu, ln :: L.LazyDataFrame +lo = lazyOf ordersDF +lc = lazyOf customersDF +lu = lazyOf usersDF +ln = lazyOf nullableDF tlo :: TL.TypedLazyDataFrame Orders tlo = TL.fromDataFrame ordersT @@ -202,5 +232,415 @@ groupAndSort = ) ] +-- | Untyped eager vs untyped lazy. +untypedPair :: String -> D.DataFrame -> L.LazyDataFrame -> Test +untypedPair lbl eager lazy = TestCase $ do + lazyDf <- L.run lazy + sameFrame (lbl ++ ": lazy == eager") eager lazyDf + +-- | Typed eager vs typed lazy. +typedPair :: + String -> DT.TypedDataFrame cols -> TL.TypedLazyDataFrame cols -> Test +typedPair lbl typed typedLazy = TestCase $ do + typedLazyDf <- TL.run typedLazy + sameFrame (lbl ++ ": typed lazy == typed") (DT.thaw typed) (DT.thaw typedLazyDf) + +-- | A value computed four ways. +valueFourWay :: (Eq a, Show a) => String -> a -> a -> IO a -> IO a -> Test +valueFourWay lbl eager typed lazy typedLazy = TestCase $ do + assertEqual (lbl ++ ": typed") eager typed + lazyV <- lazy + assertEqual (lbl ++ ": lazy") eager lazyV + typedLazyV <- typedLazy + assertEqual (lbl ++ ": typed lazy") eager typedLazyV + +g :: StdGen +g = mkStdGen 42 + +rowOps :: Test +rowOps = + TestList + [ fourWay + "filterBy" + (D.filterBy (> 4) (F.col @Double "amount") ordersDF) + (DT.filterBy (> 4) (DT.col @"amount") ordersT) + (L.filterBy (> 4) (F.col @Double "amount") lo) + (TL.filterBy (> 4) (TL.col @"amount") tlo) + , fourWay + "filterNothing" + (D.filterNothing "v" nullableDF) + (DT.filterNothing @"v" nullableT) + (L.filterNothing "v" ln) + (TL.filterNothing @"v" tln) + , fourWay + "filterAllNothing" + (D.filterAllNothing (D.exclude ["k"] nullableDF)) + (DT.filterAllNothing (DT.exclude @'["k"] nullableT)) + (L.filterAllNothing (L.exclude ["k"] ln)) + (TL.filterAllNothing (TL.exclude @'["k"] tln)) + , fourWay + "filterJust" + (D.filterJust "v" nullableDF) + (DT.filterJust @"v" nullableT) + (L.filterJust "v" ln) + (TL.filterJust @"v" tln) + , fourWay + "filterAllJust" + (D.filterAllJust nullableDF) + (DT.filterAllJust nullableT) + (L.filterAllJust ln) + (TL.filterAllJust tln) + , fourWay + "drop" + (D.drop 3 ordersDF) + (DT.drop 3 ordersT) + (L.drop 3 lo) + (TL.drop 3 tlo) + , fourWay + "takeLast" + (D.takeLast 2 ordersDF) + (DT.takeLast 2 ordersT) + (L.takeLast 2 lo) + (TL.takeLast 2 tlo) + , fourWay + "dropLast" + (D.dropLast 2 ordersDF) + (DT.dropLast 2 ordersT) + (L.dropLast 2 lo) + (TL.dropLast 2 tlo) + , fourWay + "range" + (D.range (1, 4) ordersDF) + (DT.range (1, 4) ordersT) + (L.range (1, 4) lo) + (TL.range (1, 4) tlo) + , fourWay + "range clamped" + (D.range (-2, 100) ordersDF) + (DT.range (-2, 100) ordersT) + (L.range (-2, 100) lo) + (TL.range (-2, 100) tlo) + , fourWay + "cube" + (D.cube (3, 2) ordersDF) + (DT.cube (3, 2) ordersT) + (L.cube (3, 2) lo) + (TL.cube (3, 2) tlo) + , fourWay + "topBy" + (D.topBy 1 3 [D.Desc (F.col @Double "amount")] ordersDF) + (DT.topBy 1 3 [DT.Desc (DT.col @"amount")] ordersT) + (L.topBy 1 3 [L.Desc (F.col @Double "amount")] lo) + (TL.topBy 1 3 [TL.Desc (TL.col @"amount")] tlo) + , fourWay + "distinct" + (D.distinct (D.select ["cid", "region"] ordersDF)) + (DT.distinct (DT.select @'["cid", "region"] ordersT)) + (L.distinct (L.select ["cid", "region"] lo)) + (TL.distinct (TL.select @'["cid", "region"] tlo)) + , fourWay + "sample" + (D.sample g 0.5 ordersDF) + (DT.sample g 0.5 ordersT) + (L.sample g 0.5 lo) + (TL.sample g 0.5 tlo) + , fourWay + "shuffle" + (D.shuffle g ordersDF) + (DT.shuffle g ordersT) + (L.shuffle g lo) + (TL.shuffle g tlo) + , fourWay + "selectRows" + (D.selectRows [4, 0, 2] ordersDF) + (DT.selectRows [4, 0, 2] ordersT) + (L.selectRows [4, 0, 2] lo) + (TL.selectRows [4, 0, 2] tlo) + , fourWay + "stratifiedSample" + (D.stratifiedSample g 0.5 (F.col @T.Text "region") ordersDF) + (DT.stratifiedSample g 0.5 (DT.col @"region") ordersT) + (L.stratifiedSample g 0.5 (F.col @T.Text "region") lo) + (TL.stratifiedSample g 0.5 (TL.col @"region") tlo) + , fourWay + "randomSplit fst" + (fst (D.randomSplit g 0.5 ordersDF)) + (fst (DT.randomSplit g 0.5 ordersT)) + (fst (L.randomSplit g 0.5 lo)) + (fst (TL.randomSplit g 0.5 tlo)) + , fourWay + "randomSplit snd" + (snd (D.randomSplit g 0.5 ordersDF)) + (snd (DT.randomSplit g 0.5 ordersT)) + (snd (L.randomSplit g 0.5 lo)) + (snd (TL.randomSplit g 0.5 tlo)) + , fourWay + "stratifiedSplit snd" + (snd (D.stratifiedSplit g 0.5 (F.col @T.Text "region") ordersDF)) + (snd (DT.stratifiedSplit g 0.5 (DT.col @"region") ordersT)) + (snd (L.stratifiedSplit g 0.5 (F.col @T.Text "region") lo)) + (snd (TL.stratifiedSplit g 0.5 (TL.col @"region") tlo)) + , fourWay + "kFolds last" + (last (D.kFolds g 3 ordersDF)) + (last (DT.kFolds g 3 ordersT)) + (last (L.kFolds g 3 lo)) + (last (TL.kFolds g 3 tlo)) + ] + +columnOps :: Test +columnOps = + TestList + [ fourWay + "deriveMany" + ( D.deriveMany + [ "a2" D..= (F.col @Double "amount" * 2) + , "c1" D..= (F.col @Int "cid" + 1) + ] + ordersDF + ) + ( DT.deriveMany + (DT.as @"c1" (DT.col @"cid" + 1) . DT.as @"a2" (DT.col @"amount" * 2)) + ordersT + ) + ( L.deriveMany + [ "a2" L..= (F.col @Double "amount" * 2) + , "c1" L..= (F.col @Int "cid" + 1) + ] + lo + ) + ( TL.deriveMany + (TL.as @"c1" (TL.col @"cid" + 1) . TL.as @"a2" (TL.col @"amount" * 2)) + tlo + ) + , fourWay + "deriveWithExpr" + (snd (D.deriveWithExpr "a2" (F.col @Double "amount" * 2) ordersDF)) + (snd (DT.deriveWithExpr @"a2" (DT.col @"amount" * 2) ordersT)) + (snd (L.deriveWithExpr "a2" (F.col @Double "amount" * 2) lo)) + (snd (TL.deriveWithExpr @"a2" (TL.col @"amount" * 2) tlo)) + , fourWay + "exclude" + (D.exclude ["region"] ordersDF) + (DT.exclude @'["region"] ordersT) + (L.exclude ["region"] lo) + (TL.exclude @'["region"] tlo) + , fourWay + "rename then filter on new name" + (D.filter (F.col @Double "amt" D..>. 4) (D.rename "amount" "amt" ordersDF)) + (DT.filter (DT.col @"amt" DT..>. 4) (DT.rename @"amount" @"amt" ordersT)) + (L.filter (F.col @Double "amt" L..>. 4) (L.rename "amount" "amt" lo)) + (TL.filter (TL.col @"amt" TL..>. 4) (TL.rename @"amount" @"amt" tlo)) + , fourWay + "renameMany then select" + (D.select ["c"] (D.renameMany [("cid", "c"), ("region", "r")] ordersDF)) + (DT.select @'["c"] (DT.renameMany @'[ '("cid", "c"), '("region", "r")] ordersT)) + (L.select ["c"] (L.renameMany [("cid", "c"), ("region", "r")] lo)) + (TL.select @'["c"] (TL.renameMany @'[ '("cid", "c"), '("region", "r")] tlo)) + , fourWay + "cloneColumn" + (D.cloneColumn "cid" "cid2" ordersDF) + (DT.cloneColumn @"cid" @"cid2" ordersT) + (L.cloneColumn "cid" "cid2" lo) + (TL.cloneColumn @"cid" @"cid2" tlo) + , fourWay + "dropColumn / exclude" + (D.exclude ["amount"] ordersDF) + (DT.dropColumn @"amount" ordersT) + (L.exclude ["amount"] lo) + (TL.dropColumn @"amount" tlo) + , fourWay + "insert" + (D.insert "n" [10 .. 15 :: Int] ordersDF) + (DT.insert @"n" [10 .. 15 :: Int] ordersT) + (L.insert "n" [10 .. 15 :: Int] lo) + (TL.insert @"n" [10 .. 15 :: Int] tlo) + , fourWay + "insertWithDefault" + (D.insertWithDefault 0 "n" [1, 2 :: Int] ordersDF) + (DT.insertWithDefault @"n" 0 [1, 2 :: Int] ordersT) + (L.insertWithDefault 0 "n" [1, 2 :: Int] lo) + (TL.insertWithDefault @"n" 0 [1, 2 :: Int] tlo) + , fourWay + "apply" + (D.apply (show :: Int -> String) "cid" ordersDF) + (DT.apply @"cid" (show :: Int -> String) ordersT) + (L.apply (show :: Int -> String) "cid" lo) + (TL.apply @"cid" (show :: Int -> String) tlo) + , fourWay + "applyMany" + (D.applyMany (* (2 :: Double)) ["amount"] ordersDF) + (DT.applyMany @'["amount"] (* 2) ordersT) + (L.applyMany (* (2 :: Double)) ["amount"] lo) + (TL.applyMany @'["amount"] (* 2) tlo) + , fourWay + "applyWhere" + ( D.applyWhere + (== ("us" :: T.Text)) + "region" + (+ (100 :: Double)) + "amount" + ordersDF + ) + (DT.applyWhere @"region" @"amount" (== "us") (+ 100) ordersT) + (L.applyWhere (== ("us" :: T.Text)) "region" (+ (100 :: Double)) "amount" lo) + (TL.applyWhere @"region" @"amount" (== "us") (+ 100) tlo) + , fourWay + "applyAtIndex" + (D.applyAtIndex 3 (negate :: Double -> Double) "amount" ordersDF) + (DT.applyAtIndex @"amount" 3 negate ordersT) + (L.applyAtIndex 3 (negate :: Double -> Double) "amount" lo) + (TL.applyAtIndex @"amount" 3 negate tlo) + , fourWay + "imputeWith" + (D.imputeWith F.mean (F.col @(Maybe Double) "v") nullableDF) + (DT.imputeWith DTF.mean (DT.col @"v") nullableT) + (L.imputeWith F.mean (F.col @(Maybe Double) "v") ln) + (TL.imputeWith DTF.mean (TL.col @"v") tln) + , untypedPair + "impute (untyped)" + (D.impute (F.col @(Maybe Double) "v") 0 nullableDF) + (L.impute (F.col @(Maybe Double) "v") 0 ln) + , typedPair + "impute (typed)" + (DT.impute @"v" 0 nullableT) + (TL.impute @"v" 0 tln) + , typedPair + "replaceColumn" + (DT.replaceColumn @"amount" (DT.col @"amount" * 3) ordersT) + (TL.replaceColumn @"amount" (TL.col @"amount" * 3) tlo) + ] + +mergeOps :: Test +mergeOps = + TestList + [ fourWay + "append" + (ordersDF <> ordersDF) + (DT.append ordersT ordersT) + (lo <> lo) + (TL.append tlo tlo) + , fourWay + "(|||)" + (D.select ["cid"] ordersDF D.||| D.select ["amount"] ordersDF) + (DT.select @'["cid"] ordersT DT.||| DT.select @'["amount"] ordersT) + (L.select ["cid"] lo L.||| L.select ["amount"] lo) + (TL.select @'["cid"] tlo TL.||| TL.select @'["amount"] tlo) + , fourWay + "union" + (D.union ordersDF (D.take 2 ordersDF)) + (DT.union ordersT (DT.take 2 ordersT)) + (L.union lo (L.take 2 lo)) + (TL.union tlo (TL.take 2 tlo)) + , fourWay + "intersect" + (D.intersect ordersDF (D.take 2 ordersDF)) + (DT.intersect ordersT (DT.take 2 ordersT)) + (L.intersect lo (L.take 2 lo)) + (TL.intersect tlo (TL.take 2 tlo)) + , fourWay + "difference" + (D.difference ordersDF (D.take 2 ordersDF)) + (DT.difference ordersT (DT.take 2 ordersT)) + (L.difference lo (L.take 2 lo)) + (TL.difference tlo (TL.take 2 tlo)) + , fourWay + "symmetricDifference" + (D.symmetricDifference ordersDF (D.drop 3 ordersDF)) + (DT.symmetricDifference ordersT (DT.drop 3 ordersT)) + (L.symmetricDifference lo (L.drop 3 lo)) + (TL.symmetricDifference tlo (TL.drop 3 tlo)) + , fourWay + "fold" + (D.fold (\n -> D.filter (F.col @Int "cid" D../=. F.lit n)) [1, 2] ordersDF) + (DT.fold (\n -> DT.filter (DT.col @"cid" DT../=. DT.lit n)) [1, 2] ordersT) + (L.fold (\n -> L.filter (F.col @Int "cid" L../=. F.lit n)) [1, 2] lo) + (TL.fold (\n -> TL.filter (TL.col @"cid" TL../=. TL.lit n)) [1, 2] tlo) + ] + +summaries :: Test +summaries = + TestList + [ untypedPair "summarize" (D.summarize ordersDF) (L.summarize lo) + , untypedPair "summarize (typed)" (DT.summarize ordersT) (TL.summarize tlo) + , untypedPair + "describeColumns" + (D.describeColumns ordersDF) + (L.describeColumns lo) + , untypedPair + "frequencies" + (D.frequencies (F.col @T.Text "region") ordersDF) + (L.frequencies (F.col @T.Text "region") lo) + , untypedPair + "frequencies (typed)" + (DT.frequencies (DT.col @"region") ordersT) + (TL.frequencies (TL.col @"region") tlo) + ] + +values :: Test +values = + TestList + [ valueFourWay + "mean" + (D.mean (F.col @Double "amount") ordersDF) + (DT.mean (DT.col @"amount") ordersT) + (L.mean (F.col @Double "amount") lo) + (TL.mean (TL.col @"amount") tlo) + , valueFourWay + "sum" + (D.sum (F.col @Int "cid") ordersDF) + (DT.sum (DT.col @"cid") ordersT) + (L.sum (F.col @Int "cid") lo) + (TL.sum (TL.col @"cid") tlo) + , valueFourWay + "median of a filtered query" + (D.median (F.col @Double "amount") (D.filter (F.col @Int "cid" D..>. 1) ordersDF)) + (DT.median (DT.col @"amount") (DT.filter (DT.col @"cid" DT..>. 1) ordersT)) + (L.median (F.col @Double "amount") (L.filter (F.col @Int "cid" L..>. 1) lo)) + (TL.median (TL.col @"amount") (TL.filter (TL.col @"cid" TL..>. 1) tlo)) + , valueFourWay + "dimensions" + (D.dimensions ordersDF) + (DT.dimensions ordersT) + (L.dimensions lo) + (TL.dimensions tlo) + , valueFourWay + "columnNames" + (D.columnNames ordersDF) + (DT.columnNames ordersT) + (L.columnNames lo) + (TL.columnNames tlo) + , valueFourWay + "valueCounts" + (D.valueCounts (F.col @T.Text "region") ordersDF) + (DT.valueCounts (DT.col @"region") ordersT) + (L.valueCounts (F.col @T.Text "region") lo) + (TL.valueCounts (TL.col @"region") tlo) + , valueFourWay + "columnAsList" + (D.columnAsList (F.col @Double "amount") ordersDF) + (DT.columnAsList @"amount" ordersT) + (L.columnAsList (F.col @Double "amount") lo) + (TL.columnAsList @"amount" tlo) + , valueFourWay + "correlation" + (D.correlation "cid" "amount" ordersDF) + (DT.correlation @"cid" @"amount" ordersT) + (L.correlation "cid" "amount" lo) + (TL.correlation @"cid" @"amount" tlo) + , valueFourWay + "toCsv" + (D.toCsv ordersDF) + (DT.toCsv ordersT) + (L.toCsv lo) + (TL.toCsv tlo) + , valueFourWay + "null" + (D.null ordersDF) + (DT.null ordersT) + (L.null lo) + (TL.null tlo) + ] + tests :: [Test] -tests = [joins, groupAndSort] +tests = [joins, groupAndSort, rowOps, columnOps, mergeOps, summaries, values] From 3d2315797430e08f826ee4f266eba13d39fa86bf Mon Sep 17 00:00:00 2001 From: Michael Chavinda Date: Thu, 22 Oct 2026 20:25:04 -0700 Subject: [PATCH 3/3] =?UTF-8?q?fix:=20CI=20=E2=80=94=20drop=20unused=20dat?= =?UTF-8?q?aframe-th=20dep,=20fourmolu,=20hlint=20pragmas?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../src-internal/DataFrame/Internal/Aggregation/Plan.hs | 1 - .../DataFrame/Internal/Interpreter/Kernels.hs | 1 - dataframe-lazy/src/DataFrame/Lazy/Internal/Executor.hs | 3 +-- dataframe-lazy/src/DataFrame/Typed/Lazy.hs | 1 - .../src-internal/DataFrame/DecisionTree/CondVec.hs | 1 - dataframe-operations/src/DataFrame/Functions.hs | 8 ++++++-- dataframe-operations/src/DataFrame/Typed/Expr/Extra.hs | 2 -- dataframe-operations/src/DataFrame/Typed/Join.hs | 1 - dataframe-th/dataframe-th.cabal | 1 - dataframe-th/src/DataFrame/TH/Records.hs | 2 +- 10 files changed, 8 insertions(+), 13 deletions(-) diff --git a/dataframe-core/src-internal/DataFrame/Internal/Aggregation/Plan.hs b/dataframe-core/src-internal/DataFrame/Internal/Aggregation/Plan.hs index 655ed2df..961099a6 100644 --- a/dataframe-core/src-internal/DataFrame/Internal/Aggregation/Plan.hs +++ b/dataframe-core/src-internal/DataFrame/Internal/Aggregation/Plan.hs @@ -2,7 +2,6 @@ {-# LANGUAGE ExplicitNamespaces #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE LambdaCase #-} -{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} diff --git a/dataframe-core/src-internal/DataFrame/Internal/Interpreter/Kernels.hs b/dataframe-core/src-internal/DataFrame/Internal/Interpreter/Kernels.hs index 0f820f42..fa318c98 100644 --- a/dataframe-core/src-internal/DataFrame/Internal/Interpreter/Kernels.hs +++ b/dataframe-core/src-internal/DataFrame/Internal/Interpreter/Kernels.hs @@ -1,6 +1,5 @@ {-# LANGUAGE BangPatterns #-} {-# LANGUAGE GADTs #-} -{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} diff --git a/dataframe-lazy/src/DataFrame/Lazy/Internal/Executor.hs b/dataframe-lazy/src/DataFrame/Lazy/Internal/Executor.hs index f4554d27..477e6997 100644 --- a/dataframe-lazy/src/DataFrame/Lazy/Internal/Executor.hs +++ b/dataframe-lazy/src/DataFrame/Lazy/Internal/Executor.hs @@ -560,8 +560,7 @@ inFileOrder names opts = opts{readColumns = fileOrder names <$> readColumns opts -- | Order requested names as they appear in a file (matching by leaf name). fileOrder :: [T.Text] -> [T.Text] -> [T.Text] -fileOrder names requested = - L.sortOn position requested +fileOrder names = L.sortOn position where leaf = T.takeWhileEnd (/= '.') -- Unknown names trail, so the reader still reports them as missing. diff --git a/dataframe-lazy/src/DataFrame/Typed/Lazy.hs b/dataframe-lazy/src/DataFrame/Typed/Lazy.hs index 9efae297..bfb8ee09 100644 --- a/dataframe-lazy/src/DataFrame/Typed/Lazy.hs +++ b/dataframe-lazy/src/DataFrame/Typed/Lazy.hs @@ -2,7 +2,6 @@ {-# LANGUAGE DataKinds #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE GADTs #-} -{-# LANGUAGE KindSignatures #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} {-# LANGUAGE TypeFamilies #-} diff --git a/dataframe-learn/src-internal/DataFrame/DecisionTree/CondVec.hs b/dataframe-learn/src-internal/DataFrame/DecisionTree/CondVec.hs index ff7bdf67..d785288b 100644 --- a/dataframe-learn/src-internal/DataFrame/DecisionTree/CondVec.hs +++ b/dataframe-learn/src-internal/DataFrame/DecisionTree/CondVec.hs @@ -1,5 +1,4 @@ {-# LANGUAGE GADTs #-} -{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} diff --git a/dataframe-operations/src/DataFrame/Functions.hs b/dataframe-operations/src/DataFrame/Functions.hs index 6ee75af4..7eacc446 100644 --- a/dataframe-operations/src/DataFrame/Functions.hs +++ b/dataframe-operations/src/DataFrame/Functions.hs @@ -115,7 +115,10 @@ cast colName = CastWith colName (Builtin CCast) (either (const Nothing) Just) -- | Lenient coercion that substitutes a default for unconvertible values. castWithDefault :: forall a. (Columnable a, Read a) => a -> T.Text -> Expr a castWithDefault def colName = - CastWith colName (Builtin (CCastWithDefault (T.pack (show def)))) (fromRight def) + CastWith + colName + (Builtin (CCastWithDefault (T.pack (show def)))) + (fromRight def) -- | Lenient coercion to @Either T.Text a@, keeping unparsed values as 'Left'. castEither :: @@ -342,7 +345,8 @@ bottomK = kExtremes False {-# INLINEABLE bottomK #-} -- | The k largest (True) or smallest (False) values, dropping NaNs. -kExtremes :: forall a. (Columnable a, Ord a) => Bool -> Int -> Expr a -> Expr [a] +kExtremes :: + forall a. (Columnable a, Ord a) => Bool -> Int -> Expr a -> Expr [a] kExtremes largest k = Agg ( MergeAgg diff --git a/dataframe-operations/src/DataFrame/Typed/Expr/Extra.hs b/dataframe-operations/src/DataFrame/Typed/Expr/Extra.hs index bcb2244f..f5d182ce 100644 --- a/dataframe-operations/src/DataFrame/Typed/Expr/Extra.hs +++ b/dataframe-operations/src/DataFrame/Typed/Expr/Extra.hs @@ -1,8 +1,6 @@ {-# LANGUAGE AllowAmbiguousTypes #-} {-# LANGUAGE DataKinds #-} {-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE KindSignatures #-} -{-# LANGUAGE MonoLocalBinds #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} {-# LANGUAGE TypeFamilies #-} diff --git a/dataframe-operations/src/DataFrame/Typed/Join.hs b/dataframe-operations/src/DataFrame/Typed/Join.hs index cec60e56..b300953c 100644 --- a/dataframe-operations/src/DataFrame/Typed/Join.hs +++ b/dataframe-operations/src/DataFrame/Typed/Join.hs @@ -3,7 +3,6 @@ {-# LANGUAGE DataKinds #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE GADTs #-} -{-# LANGUAGE KindSignatures #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} diff --git a/dataframe-th/dataframe-th.cabal b/dataframe-th/dataframe-th.cabal index 860ae9ea..c2fdbe2b 100644 --- a/dataframe-th/dataframe-th.cabal +++ b/dataframe-th/dataframe-th.cabal @@ -36,7 +36,6 @@ library build-depends: base >= 4 && < 5, containers >= 0.6.7 && < 0.10, dataframe-core >= 2.6 && < 2.7, - dataframe-operations >= 2.6 && < 2.7, dataframe-parsing >= 2.2 && < 2.3, template-haskell >= 2.0 && < 3, text >= 2.1 && < 3, diff --git a/dataframe-th/src/DataFrame/TH/Records.hs b/dataframe-th/src/DataFrame/TH/Records.hs index c9b01b23..0ec05127 100644 --- a/dataframe-th/src/DataFrame/TH/Records.hs +++ b/dataframe-th/src/DataFrame/TH/Records.hs @@ -48,8 +48,8 @@ import DataFrame.Internal.DataFrame ( ) import DataFrame.Internal.Expression (Expr) import DataFrame.Internal.Schema.TH (schemaTypeE) -import DataFrame.TH.Identifier (sanitize) import DataFrame.Schema (Schema, makeSchema) +import DataFrame.TH.Identifier (sanitize) import Prelude as P knownTypeNames :: M.Map String Name