diff --git a/dataframe-lazy.cabal b/dataframe-lazy.cabal --- a/dataframe-lazy.cabal +++ b/dataframe-lazy.cabal @@ -1,6 +1,6 @@ cabal-version: 3.4 name: dataframe-lazy -version: 2.4.2.1 +version: 2.5.0.0 synopsis: Lazy query engine for the dataframe ecosystem. description: The lazy/streaming query engine: relational-algebra plans, optimizer, @@ -36,6 +36,7 @@ -- 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 @@ -45,14 +46,16 @@ attoparsec >= 0.12 && < 0.16, bytestring >= 0.11 && < 0.14, containers >= 0.6.7 && < 0.10, - dataframe-core >= 2.6 && < 2.7, + dataframe-core >= 2.7 && < 2.8, dataframe-csv >= 2.3 && < 2.4, - dataframe-operations >= 2.6 && < 2.7, + dataframe-operations >= 2.7 && < 2.8, dataframe-parquet >= 1.6 && < 1.7, dataframe-parsing >= 2.2 && < 2.3, + dataframe-viz >= 1.4 && < 1.5, 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, @@ -68,7 +71,7 @@ build-depends: base >= 4 && < 5, bytestring >= 0.11 && < 0.14, containers >= 0.6.7 && < 0.10, - dataframe-core >= 2.6 && < 2.7, + dataframe-core >= 2.7 && < 2.8, dataframe-csv >= 2.3 && < 2.4, dataframe-lazy, dataframe-parsing >= 2.2 && < 2.3, diff --git a/src/DataFrame/Lazy.hs b/src/DataFrame/Lazy.hs --- a/src/DataFrame/Lazy.hs +++ b/src/DataFrame/Lazy.hs @@ -1,7 +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.Lazy.Internal.LogicalPlan (SortOrder (..)) +import DataFrame.Operations.Join (JoinType (..)) +import DataFrame.Operations.Permutation (SortOrder (..)) diff --git a/src/DataFrame/Lazy/Internal/Collect.hs b/src/DataFrame/Lazy/Internal/Collect.hs new file mode 100644 --- /dev/null +++ b/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/src/DataFrame/Lazy/Internal/DataFrame.hs b/src/DataFrame/Lazy/Internal/DataFrame.hs --- a/src/DataFrame/Lazy/Internal/DataFrame.hs +++ b/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,17 +19,35 @@ 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 (..), - SortOrder (..), + opaque, ) import qualified DataFrame.Lazy.Internal.Optimizer as Opt -import DataFrame.Operations.Join (JoinType) +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 'runDataFrame'. +-- | A lazy query: a 'LogicalPlan' whose execution waits for 'run'. data LazyDataFrame = LazyDataFrame { plan :: LogicalPlan , batchSize :: Int @@ -39,8 +59,8 @@ <> (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,38 +133,383 @@ 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 --- | Sort the result by the given @(column, direction)@ pairs. -sortBy :: [(T.Text, SortOrder)] -> LazyDataFrame -> LazyDataFrame -sortBy cols ldf = ldf{plan = Sort cols (plan ldf)} +-- | 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 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 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/src/DataFrame/Lazy/Internal/Executor.hs b/src/DataFrame/Lazy/Internal/Executor.hs --- a/src/DataFrame/Lazy/Internal/Executor.hs +++ b/src/DataFrame/Lazy/Internal/Executor.hs @@ -25,16 +25,14 @@ 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.List as L 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, @@ -42,11 +40,17 @@ 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 (..), SortOrder (..)) +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 @@ -75,6 +79,10 @@ 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 [] @@ -109,8 +117,13 @@ 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 (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 -- | Execute a physical plan into a single complete 'DataFrame'. execute :: PhysicalPlan -> IO D.DataFrame @@ -149,7 +162,7 @@ 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 @@ -165,6 +178,32 @@ 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) @@ -182,10 +221,9 @@ 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 +268,26 @@ 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 +295,7 @@ 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 +312,10 @@ 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 :: @@ -351,6 +386,13 @@ where combine (p1, m1, f1) (p2, m2, f2) = (p1 ++ p2, m1 ++ m2, f1 . f2) + -- Partial and merge steps are new operators, never the builtin itself. + phase :: T.Text -> E.OpRef E.AggId -> E.OpRef E.AggId + phase p ident = case ident of + E.Builtin a -> E.Named (p <> E.aggName a) + E.Named n -> E.Named (p <> n) + E.Anonymous l -> E.Anonymous (p <> l) + processAgg :: (T.Text, E.SomeExpr) -> ([(T.Text, E.SomeExpr)], [(T.Text, E.SomeExpr)], D.DataFrame -> D.DataFrame) @@ -370,7 +412,10 @@ [ ( name , E.SomeExpr - (E.Agg (E.FoldAgg "sum" Nothing ((+) :: Int -> Int -> Int)) (E.Col @Int name)) + ( E.Agg + (E.FoldAgg (E.Builtin E.ASum) Nothing ((+) :: Int -> Int -> Int)) + (E.Col @Int name) + ) ) ] , id @@ -384,7 +429,10 @@ [ ( name , E.SomeExpr - (E.Agg (E.FoldAgg "sum" Nothing ((+) :: Int -> Int -> Int)) (E.Col @Int name)) + ( E.Agg + (E.FoldAgg (E.Builtin E.ASum) Nothing ((+) :: Int -> Int -> Int)) + (E.Col @Int name) + ) ) ] , id @@ -411,13 +459,13 @@ let partialExpr = E.SomeExpr ( E.Agg - (E.FoldAgg ("partial_" <> n) (Just seed) step) + (E.FoldAgg (phase "partial_" n) (Just seed) step) inner ) mergeExpr = E.SomeExpr ( E.Agg - (E.FoldAgg ("merge_" <> n) Nothing merge) + (E.FoldAgg (phase "merge_" n) Nothing merge) (E.Col @acc name) ) finalize df = @@ -455,7 +503,12 @@ [] -> 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 @@ -466,9 +519,10 @@ 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 @@ -487,10 +541,34 @@ 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') ) +-- | 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 = L.sortOn position + 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 = @@ -508,7 +586,7 @@ 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 @@ -517,6 +595,7 @@ ( withFile path ReadMode $ \h -> do rawHeader <- C8.hGetLine h let header = stripUtf8Bom rawHeader + opts = inFileOrder (headerNames sep header) schemaOpts let feed bytes = unless (BS.null bytes) $ do let window = header <> BS.singleton nl <> bytes @@ -555,7 +634,7 @@ 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 @@ -628,61 +707,17 @@ 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) + 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) - 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 +-- | 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/src/DataFrame/Lazy/Internal/LogicalPlan.hs b/src/DataFrame/Lazy/Internal/LogicalPlan.hs --- a/src/DataFrame/Lazy/Internal/LogicalPlan.hs +++ b/src/DataFrame/Lazy/Internal/LogicalPlan.hs @@ -7,6 +7,7 @@ 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,29 @@ "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) +-- | 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. @@ -40,14 +60,24 @@ 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 + | -- | 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/src/DataFrame/Lazy/Internal/Optimizer.hs b/src/DataFrame/Lazy/Internal/Optimizer.hs --- a/src/DataFrame/Lazy/Internal/Optimizer.hs +++ b/src/DataFrame/Lazy/Internal/Optimizer.hs @@ -3,12 +3,14 @@ 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 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,13 +28,13 @@ 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) 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 @@ -40,7 +42,7 @@ E.Binary ( E.MkBinaryOp { E.binaryFn = (&&) - , E.binaryName = "and" + , E.binaryId = E.Builtin E.BAnd , E.binarySymbol = Just "&&" , E.binaryCommutative = True , E.binaryPrecedence = 3 @@ -59,17 +61,43 @@ 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) -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) 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) referencedCols (Scan _ schema) = Just (S.fromList (M.keys (elements schema))) @@ -78,19 +106,34 @@ 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 (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 @@ -112,16 +155,28 @@ 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 (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 @@ -129,6 +184,14 @@ 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. @@ -170,17 +233,21 @@ 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 (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/src/DataFrame/Lazy/Internal/PhysicalPlan.hs b/src/DataFrame/Lazy/Internal/PhysicalPlan.hs --- a/src/DataFrame/Lazy/Internal/PhysicalPlan.hs +++ b/src/DataFrame/Lazy/Internal/PhysicalPlan.hs @@ -3,8 +3,9 @@ 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 (BatchFn, BatchFn2, DataSource) import DataFrame.Operations.Join (JoinType) +import DataFrame.Operations.Permutation (SortOrder) import DataFrame.Schema (Schema) -- | Scan-level configuration: batch size, separator, optional pushdowns. @@ -22,11 +23,16 @@ | 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 + | 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/src/DataFrame/Typed/Lazy.hs b/src/DataFrame/Typed/Lazy.hs --- a/src/DataFrame/Typed/Lazy.hs +++ b/src/DataFrame/Typed/Lazy.hs @@ -2,236 +2,1021 @@ {-# LANGUAGE DataKinds #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE GADTs #-} -{-# LANGUAGE KindSignatures #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} - --- | Type-safe lazy query pipelines with compile-time schema tracking. -module DataFrame.Typed.Lazy ( - -- * Core type - TypedLazyDataFrame, - - -- * Data sources - scanCsv, - scanSeparated, - scanParquet, - fromDataFrame, - fromTypedDataFrame, - - -- * Schema-preserving operations - filter, - take, - - -- * Schema-modifying operations - derive, - select, - - -- * Aggregation - groupBy, - aggregate, - - -- * Joins - innerJoin, - leftJoin, - rightJoin, - fullOuterJoin, - - -- * Sort - sortBy, - - -- * Execution - run, - - -- * Re-exports for pipeline construction - module DataFrame.Typed.Expr, - module DataFrame.Typed.Types, - SortOrder (..), -) where - -import Data.Kind (Type) -import Data.Proxy (Proxy (..)) -import qualified Data.Text as T -import GHC.TypeLits (KnownSymbol, Symbol, symbolVal) -import Prelude hiding (filter, take) - -import qualified DataFrame.Internal.Column as C -import qualified DataFrame.Internal.Expression as E -import DataFrame.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.Typed.Expr -import DataFrame.Typed.Freeze (unsafeFreeze) -import DataFrame.Typed.Schema -import DataFrame.Typed.Types - --- | A lazy query with compile-time schema tracking. -newtype TypedLazyDataFrame (cols :: [(Symbol, Type)]) = TLD {_unTLD :: LazyDataFrame} - -instance Show (TypedLazyDataFrame cols) where - show (TLD ldf) = "TypedLazyDataFrame { " ++ show ldf ++ " }" - --- | Scan a CSV file with a given schema. -scanCsv :: - Schema -> - T.Text -> - TypedLazyDataFrame cols -scanCsv schema path = TLD (L.scanCsv schema path) - --- | Scan a character-separated file with a given schema. -scanSeparated :: - Char -> - Schema -> - T.Text -> - TypedLazyDataFrame cols -scanSeparated sep schema path = TLD (L.scanSeparated sep schema path) - --- | Scan a Parquet file, directory, or glob pattern with a given schema. -scanParquet :: - Schema -> - T.Text -> - TypedLazyDataFrame cols -scanParquet schema path = TLD (L.scanParquet schema path) - --- | Lift an already-loaded eager 'TypedDataFrame' into a lazy plan. -fromDataFrame :: TypedDataFrame cols -> TypedLazyDataFrame cols -fromDataFrame (TDF df) = TLD (L.fromDataFrame df) - --- | Synonym for 'fromDataFrame'. -fromTypedDataFrame :: TypedDataFrame cols -> TypedLazyDataFrame cols -fromTypedDataFrame = fromDataFrame - --- | Keep rows that satisfy the predicate. -filter :: TExpr cols Bool -> TypedLazyDataFrame cols -> TypedLazyDataFrame cols -filter (TExpr expr) (TLD ldf) = TLD (L.filter expr ldf) - --- | Retain at most @n@ rows. -take :: Int -> TypedLazyDataFrame cols -> TypedLazyDataFrame cols -take n (TLD ldf) = TLD (L.take n ldf) - --- | Add a computed column. -derive :: - forall name a cols. - (KnownSymbol name, C.Columnable a, AssertAbsent name cols) => - TExpr cols a -> - TypedLazyDataFrame cols -> - TypedLazyDataFrame (Snoc cols '(name, a)) -derive (TExpr expr) (TLD ldf) = - TLD (L.derive (T.pack (symbolVal (Proxy @name))) expr ldf) - --- | Retain only the listed columns. -select :: - forall (names :: [Symbol]) cols. - (AllKnownSymbol names, AssertAllPresent names cols) => - TypedLazyDataFrame cols -> - TypedLazyDataFrame (SubsetSchema names cols) -select (TLD ldf) = TLD (L.select (DataFrame.Typed.Schema.symbolVals @names) ldf) - --- | A typed lazy grouped query. -newtype TypedLazyGrouped (keys :: [Symbol]) (cols :: [(Symbol, Type)]) = TLG - { _unTLG :: ([T.Text], LazyDataFrame) - } - --- | Group by key columns. -groupBy :: - forall (keys :: [Symbol]) cols. - (AllKnownSymbol keys, AssertAllPresent keys cols) => - TypedLazyDataFrame cols -> - TypedLazyGrouped keys cols -groupBy (TLD ldf) = TLG (DataFrame.Typed.Schema.symbolVals @keys, ldf) - --- | Aggregate a grouped lazy query. -aggregate :: - forall keys cols aggs. - TAgg keys cols aggs -> - TypedLazyGrouped keys cols -> - TypedLazyDataFrame (Append (GroupKeyColumns keys cols) (Reverse aggs)) -aggregate tagg (TLG (keys, ldf)) = - TLD (L.groupBy keys (aggToNamedExprs tagg) ldf) - --- | 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 - ) => - TypedLazyDataFrame left -> - TypedLazyDataFrame right -> - TypedLazyDataFrame (InnerJoinSchema '[key] left right) -innerJoin = joinOn @key INNER - --- | 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 - ) => - TypedLazyDataFrame left -> - TypedLazyDataFrame right -> - TypedLazyDataFrame (LeftJoinSchema '[key] left right) -leftJoin = joinOn @key LEFT - --- | 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 - ) => - TypedLazyDataFrame left -> - TypedLazyDataFrame right -> - TypedLazyDataFrame (RightJoinSchema '[key] left right) -rightJoin = joinOn @key RIGHT - --- | 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 - ) => - TypedLazyDataFrame left -> - TypedLazyDataFrame right -> - TypedLazyDataFrame (FullOuterJoinSchema '[key] left right) -fullOuterJoin = joinOn @key FULL_OUTER - --- | Runtime delegation shared by the typed joins. -joinOn :: - forall (key :: Symbol) left right out. - (KnownSymbol key) => - JoinType -> - 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)) - --- | Sort the result by column name and direction. -sortBy :: - [(T.Text, SortOrder)] -> - TypedLazyDataFrame cols -> - TypedLazyDataFrame cols -sortBy cols (TLD ldf) = TLD (L.sortBy cols ldf) - --- | Execute the lazy query and return a typed DataFrame. -run :: - forall cols. - (KnownSchema cols) => - TypedLazyDataFrame cols -> - IO (TypedDataFrame cols) -run (TLD ldf) = unsafeFreeze <$> L.runDataFrame ldf - --- | Convert TAgg to untyped named expressions for the lazy groupBy. -aggToNamedExprs :: TAgg keys cols aggs -> [(T.Text, E.SomeExpr)] -aggToNamedExprs TAggNil = [] -aggToNamedExprs (TAggCons name (TExpr expr) rest) = - (name, E.SomeExpr expr) : aggToNamedExprs rest +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} + +-- | Typed lazy queries with the same operations as "DataFrame.Typed". +module DataFrame.Typed.Lazy ( + -- * Core types + TypedLazyDataFrame, + TypedLazyGrouped, + TypedDataFrame, + + -- * Data sources + scanCsv, + scanCsvWith, + scanSeparated, + scanSeparatedWith, + scanParquet, + fromDataFrame, + + -- * 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, + as, + aggregate, + + -- * Joins + JoinType (..), + join, + joinOn, + innerJoin, + leftJoin, + rightJoin, + fullOuterJoin, + + -- * 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 (..), + col, + lit, + ifThenElse, + caseWhen, + orElse, + (.==.), + (./=.), + (.<.), + (.<=.), + (.>=.), + (.>.), + (.+.), + (.-.), + (.*.), + (./.), + (.^^.), + (.^.), + (.+), + (.-), + (.*), + (./), + (.^^), + (.^), + (.==), + (./=), + (.<), + (.<=), + (.>=), + (.>), + (.&&.), + (.||.), + (.&&), + (.||), + (|>), +) where + +import Data.Kind (Type) +import Data.Proxy (Proxy (..)) +import qualified Data.Text as T +import GHC.TypeLits (KnownSymbol, Symbol, symbolVal) +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 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, + JoinOnSchema, + JoinSchema, + KnownJoinType (..), + 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} + +instance Show (TypedLazyDataFrame cols) where + show (TLD ldf) = "TypedLazyDataFrame { " ++ show ldf ++ " }" + +-- | Scan a CSV file; the column types come from @cols@. +scanCsv :: + 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; the column types come from @cols@. +scanSeparated :: + 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; the columns come from @cols@. +scanParquet :: + 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) + +-- | Keep rows that satisfy the predicate. +filter :: TExpr cols Bool -> TypedLazyDataFrame cols -> TypedLazyDataFrame cols +filter (TExpr expr) (TLD ldf) = TLD (L.filter expr ldf) + +-- | Retain at most @n@ rows. +take :: Int -> TypedLazyDataFrame cols -> TypedLazyDataFrame cols +take n (TLD ldf) = TLD (L.take n ldf) + +-- | 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. + (KnownSymbol name, C.Columnable a, AssertAbsent name cols) => + TExpr cols a -> + TypedLazyDataFrame cols -> + TypedLazyDataFrame (Snoc cols '(name, a)) +derive (TExpr expr) (TLD ldf) = + TLD (L.derive (T.pack (symbolVal (Proxy @name))) expr ldf) + +-- | Retain only the listed columns. +select :: + forall (names :: [Symbol]) cols. + (AllKnownSymbol names, AssertAllPresent names cols) => + TypedLazyDataFrame cols -> + TypedLazyDataFrame (SubsetSchema names cols) +select (TLD ldf) = TLD (L.select (symbolVals @names) ldf) + +-- | A typed lazy query grouped by @keys@, awaiting 'aggregate'. +newtype TypedLazyGrouped (keys :: [Symbol]) (cols :: [(Symbol, Type)]) + = TLG L.LazyGroupedDataFrame + +-- | 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 (L.groupBy (symbolVals @keys) ldf) + +-- | Aggregate each group; build the outputs with 'as', as in the eager API. +aggregate :: + forall keys cols aggs. + (TAgg keys cols '[] -> TAgg keys cols aggs) -> + TypedLazyGrouped keys cols -> + TypedLazyDataFrame (Append (GroupKeyColumns keys cols) (Reverse aggs)) +aggregate build (TLG grouped) = + TLD (L.aggregate (taggToNamedExprs (build TAggNil)) grouped) + +-- | 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 (JoinSchema jt keys left right) +join (TLD l) (TLD r) = TLD (L.join (joinTypeVal @jt) (symbolVals @keys) l r) + +-- | 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 (JoinOnSchema jt pairs left right) +joinOn (TLD l) (TLD r) = TLD (L.joinOn (joinTypeVal @jt) (pairVals @pairs) l r) + +-- | Typed inner join on one or more key columns. +innerJoin :: + forall (keys :: [Symbol]) left right. + (JoinKeys keys left right) => + TypedLazyDataFrame left -> + TypedLazyDataFrame right -> + TypedLazyDataFrame (InnerJoinSchema keys left right) +innerJoin = join @'INNER @keys + +-- | 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 (LeftJoinSchema keys left right) +leftJoin = join @'LEFT @keys + +-- | 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 (RightJoinSchema keys left right) +rightJoin = join @'RIGHT @keys + +-- | 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 :: 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/test/StreamingCsv.hs b/test/StreamingCsv.hs --- a/test/StreamingCsv.hs +++ b/test/StreamingCsv.hs @@ -46,7 +46,7 @@ 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 @@ 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 @@ 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 @@ 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 @@ withRawCsv input $ \path -> do result <- try - ( L.runDataFrame + ( L.run (L.scanCsvStreamingBytesWith failingReader schema (T.pack path)) ) case result of @@ -155,7 +155,7 @@ 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 ()