From eb332ed1d583054ce8463ccb9aae8e650800349c Mon Sep 17 00:00:00 2001 From: Michael Chavinda Date: Thu, 22 Oct 2026 16:47:49 -0700 Subject: [PATCH] feat: Durable operation identity String matching kinda sucked and I meant to do this eventualyl anyway. Would have been tedious. Had Claude do it. --- CHANGELOG.md | 41 +++ dataframe-core/dataframe-core.cabal | 1 + .../DataFrame/Internal/Aggregation/Plan.hs | 48 ++-- .../DataFrame/Internal/Expression.hs | 219 ++++++++++----- .../DataFrame/Internal/Expression/OpId.hs | 260 ++++++++++++++++++ .../Internal/Expression/Operators.hs | 75 +++-- .../DataFrame/Internal/Expression/Simplify.hs | 24 +- .../DataFrame/Internal/Interpreter.hs | 82 +++--- .../DataFrame/Internal/Interpreter/Kernels.hs | 31 ++- .../src/DataFrame/IR/ExprJson.hs | 164 ++++++----- .../src/DataFrame/Lazy/Internal/Executor.hs | 21 +- .../src/DataFrame/Lazy/Internal/Optimizer.hs | 2 +- .../DataFrame/DecisionTree/CondVec.hs | 49 ++-- .../DataFrame/DecisionTree/Pool.hs | 34 ++- .../DataFrame/Internal/Statistics.hs | 4 + .../src/DataFrame/Functions.hs | 229 ++++++--------- .../src/DataFrame/Operations/Statistics.hs | 7 - .../src/DataFrame/Typed/Expr.hs | 124 +++++---- .../src/DataFrame/Typed/Expr/Extra.hs | 127 ++++++++- dataframe-parquet/src/DataFrame/IO/Parquet.hs | 35 +-- dataframe-th/dataframe-th.cabal | 1 + dataframe-th/src/DataFrame/TH/Identifier.hs | 75 +++++ dataframe-th/src/DataFrame/TH/Records.hs | 2 +- .../DataFrame/Display/Internal/VegaLite.hs | 9 +- .../src/DataFrame/Display/Web/Chart/Typed.hs | 5 +- dataframe.cabal | 2 + examples/examples.cabal | 2 + src/DataFrame/Typed.hs | 18 ++ tests/Functions.hs | 4 +- tests/IR/ExprJsonRoundtrip.hs | 6 + tests/Main.hs | 2 + tests/Operations/OperatorIdentity.hs | 255 +++++++++++++++++ tests/Typed/Parity.hs | 37 +++ 33 files changed, 1465 insertions(+), 530 deletions(-) create mode 100644 dataframe-core/src-internal/DataFrame/Internal/Expression/OpId.hs create mode 100644 dataframe-th/src/DataFrame/TH/Identifier.hs create mode 100644 tests/Operations/OperatorIdentity.hs diff --git a/CHANGELOG.md b/CHANGELOG.md index a400193c..92e11289 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,46 @@ # Revision history for dataframe +## Unreleased + +### Breaking changes +* Operators are identified by tag, not name. `UnaryOp`/`BinaryOp` expose + `unaryId`/`binaryId :: OpRef …` (`Builtin` tag, `Named` text or + `Anonymous` label); `unaryName`/`binaryName` are now plain functions. + `MkUnaryOp`/`MkBinaryOp` take an `OpRef` instead of a name, and + `AggStrategy` constructors take an `OpRef AggId`. +* `liftDecorated`/`lift2Decorated` build `Named` ops, which never use a + builtin kernel; library code uses the new `liftBuiltin`/`lift2Builtin`. +* `CastWith`/`CastExprWith` take an `OpRef CastId` instead of a text tag. +* `kExtremes` takes a `Bool` (largest/smallest) instead of a tag and a + comparator, so the two can no longer disagree. +* Every `DataFrame.Functions` combinator that does not take a Haskell + function is a builtin; `lift`, `lift2`, `nullLift`, `nullLift2`, + `reduce`, `whenPresent`, `whenBothPresent` and `bind` stay anonymous. +* Typed `not` prints as `~`, like the untyped `not`. +* `DataFrame.Functions` holds only expression combinators: `sanitize`, + `isHaskellIdentifier`, `isReservedId` and `isVarId` move to + `DataFrame.TH.Identifier` (dataframe-th), and `optionalToDoubleVector` to + `DataFrame.Internal.Statistics`. + +### Additions +* `DataFrame.Typed` gains `cast`, `castWithDefault`, `castEither`, + `unsafeCast` (schema-checked column), `eq`, `lt`, `gt`, `leq`, `geq`, + `and`, `or`, `countDistinct`, `topK`, `bottomK`, `isNull`, `isNotNull`, + `contains` and `length`, completing parity with `DataFrame.Functions`. + +### Fixes +* A user function named like a builtin (e.g. `"add"`) no longer runs the + builtin's kernel instead of the user's function. +* `eqExpr` no longer equates different `lift`/`lift2` functions, `fromMaybe` + with different defaults, or `(^)` with `(^^)`. +* The predicate simplifier no longer merges conditions on different + user functions, and no longer treats `fromMaybe True` as false on nulls. +* The grouped moments fast path no longer drops a user function inside `sum`. +* `medianMaybe` prints as `medianMaybe`, not `meanMaybe`. +* Parquet row-group pruning recognises `F.and`/`.&&` conjunctions. +* The expression wire format decodes `logBase`, `stddev` and + `countDistinct`; `(^^)` encodes as `powFrac`. + ## 3.7.0.0 * `dataframe-viz`: column names are escaped in Vega-Lite field references diff --git a/dataframe-core/dataframe-core.cabal b/dataframe-core/dataframe-core.cabal index bdb5de1c..b769795d 100644 --- a/dataframe-core/dataframe-core.cabal +++ b/dataframe-core/dataframe-core.cabal @@ -76,6 +76,7 @@ library -- Expressions: syntax, operators, simplification. DataFrame.Internal.Expression DataFrame.Internal.Expression.Operators + DataFrame.Internal.Expression.OpId DataFrame.Internal.Expression.Reflect DataFrame.Internal.Expression.Operators.Nullable DataFrame.Internal.Expression.Simplify diff --git a/dataframe-core/src-internal/DataFrame/Internal/Aggregation/Plan.hs b/dataframe-core/src-internal/DataFrame/Internal/Aggregation/Plan.hs index f3d6a176..655ed2df 100644 --- a/dataframe-core/src-internal/DataFrame/Internal/Aggregation/Plan.hs +++ b/dataframe-core/src-internal/DataFrame/Internal/Aggregation/Plan.hs @@ -28,10 +28,15 @@ import DataFrame.Internal.DataFrame ( getColumn, ) import DataFrame.Internal.Expression ( + AggId (..), AggStrategy (..), - BinaryOp (binaryCommutative, binaryName), + BinOpId (..), + BinaryOp (binaryId), Expr (..), + OpRef (..), SomeExpr (..), + UnOpId (..), + UnaryOp (unaryId), ) import Type.Reflection (Typeable, typeRep) @@ -49,27 +54,25 @@ data AggPlan -- | Plan a named output expression, or 'Nothing' if the shape is unsupported. planAgg :: GroupedDataFrame -> SomeExpr -> Maybe AggPlan planAgg gdf (SomeExpr (expr :: Expr a)) = case expr of - Agg (FoldAgg tag _ _) (Col name) -> foldPlan tag name - Agg (MergeAgg tag _ _ _ _) (Col name) -> mergePlan tag name - Agg (CollectAgg tag _) (Col name) -> collectPlan tag name + Agg (FoldAgg (Builtin tag) _ _) (Col name) -> foldPlan tag name + Agg (MergeAgg (Builtin tag) _ _ _ _) (Col name) -> mergePlan tag name + Agg (CollectAgg (Builtin tag) _) (Col name) -> collectPlan tag name Binary op - (Agg (FoldAgg lt Nothing _) (Col a)) - (Agg (FoldAgg rt Nothing _) (Col b)) -> - if binaryName op == "sub" && lt == "maximum" && rt == "minimum" - then requireBoth a b (PlanMaxMinusMin a b) - else Nothing + (Agg (FoldAgg (Builtin AMaximum) Nothing _) (Col a)) + (Agg (FoldAgg (Builtin AMinimum) Nothing _) (Col b)) + | Builtin BSub <- binaryId op -> requireBoth a b (PlanMaxMinusMin a b) _ -> Nothing where foldPlan tag name = case tag of - "sum" -> require name (PlanScatter RSum name) - "minimum" -> require name (PlanScatter RMin name) - "maximum" -> require name (PlanScatter RMax name) + ASum -> require name (PlanScatter RSum name) + AMinimum -> require name (PlanScatter RMin name) + AMaximum -> require name (PlanScatter RMax name) _ -> Nothing mergePlan tag name = case tag of - "mean" -> outputType @Double >> require name (PlanScatter RMean name) - "count" -> outputType @Int >> require name (PlanScatter RCount name) - "countDistinct" -> do + AMean -> outputType @Double >> require name (PlanScatter RMean name) + ACount -> outputType @Int >> require name (PlanScatter RCount name) + ACountDistinct -> do outputType @Int c <- getColumn name (fullDataframe gdf) if distinctSupported c then Just (PlanCountDistinct name) else Nothing @@ -79,9 +82,9 @@ planAgg gdf (SomeExpr (expr :: Expr a)) = case expr of Just Refl -> Just () Nothing -> Nothing collectPlan tag name = case tag of - "stddev" -> require name (PlanScatter RStd name) - "variance" -> require name (PlanScatter RVar name) - "median" -> require name (PlanMedian name) + AStddev -> require name (PlanScatter RStd name) + AVariance -> require name (PlanScatter RVar name) + AMedian -> require name (PlanMedian name) _ -> Nothing require name plan = colUnboxedNumeric name >> Just plan requireBoth a b plan = colUnboxedNumeric a >> colUnboxedNumeric b >> Just plan @@ -162,8 +165,8 @@ data Role -- | Tag a single named aggregation with its moment role, or reject the group. classify :: M.Map T.Text SomeExpr -> (T.Text, SomeExpr) -> Maybe (T.Text, Role) classify exprs (name, SomeExpr expr) = case expr of - Agg (MergeAgg "count" _ _ _ _) _ -> Just (name, RoleN) - Agg (FoldAgg "sum" _ _) arg -> (\t -> (name, termRole t)) <$> resolveTerm exprs (SomeExpr arg) + Agg (MergeAgg (Builtin ACount) _ _ _ _) _ -> Just (name, RoleN) + Agg (FoldAgg (Builtin ASum) _ _) arg -> (\t -> (name, termRole t)) <$> resolveTerm exprs (SomeExpr arg) _ -> Nothing termRole :: Term -> Role @@ -179,9 +182,10 @@ resolveTerm exprs = go (8 :: Int) Col nm -> case M.lookup nm exprs of Just ue -> go (fuel - 1) ue Nothing -> Just (Lin nm) - Unary _ inner -> go (fuel - 1) (SomeExpr inner) + Unary op inner + | Builtin UToDouble <- unaryId op -> go (fuel - 1) (SomeExpr inner) Binary op l r - | binaryName op == "mult" && binaryCommutative op -> do + | Builtin BMult <- binaryId op -> do Lin a <- go (fuel - 1) (SomeExpr l) Lin b <- go (fuel - 1) (SomeExpr r) Just (sortProd a b) diff --git a/dataframe-core/src-internal/DataFrame/Internal/Expression.hs b/dataframe-core/src-internal/DataFrame/Internal/Expression.hs index 3ac29c9b..b1145c36 100644 --- a/dataframe-core/src-internal/DataFrame/Internal/Expression.hs +++ b/dataframe-core/src-internal/DataFrame/Internal/Expression.hs @@ -13,7 +13,10 @@ {-# LANGUAGE UndecidableInstances #-} {-# LANGUAGE NoFieldSelectors #-} -module DataFrame.Internal.Expression where +module DataFrame.Internal.Expression ( + module DataFrame.Internal.Expression, + module DataFrame.Internal.Expression.OpId, +) where import qualified Data.Map.Strict as M import Data.Maybe (fromMaybe) @@ -23,18 +26,19 @@ import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl)) import qualified Data.Vector.Generic as VG import DataFrame.Internal.Column import qualified DataFrame.Internal.Display.Pretty as P +import DataFrame.Internal.Expression.OpId import Type.Reflection (Typeable, typeOf, typeRep) -- | Operators are an open typeclass so the simplifier can match them by 'cast'. class (Typeable op) => UnaryOp op where unaryFn :: op a b -> a -> b - unaryName :: op a b -> T.Text + unaryId :: op a b -> OpRef UnOpId unarySymbol :: op a b -> Maybe T.Text unarySymbol _ = Nothing class (Typeable op) => BinaryOp op where binaryFn :: op a b c -> a -> b -> c - binaryName :: op a b c -> T.Text + binaryId :: op a b c -> OpRef BinOpId binarySymbol :: op a b c -> Maybe T.Text binarySymbol _ = Nothing binaryCommutative :: op a b c -> Bool @@ -42,15 +46,23 @@ class (Typeable op) => BinaryOp op where binaryPrecedence :: op a b c -> Int binaryPrecedence _ = 9 +-- | Display name of a unary operator. +unaryName :: (UnaryOp op) => op a b -> T.Text +unaryName = opRefName unOpName . unaryId + +-- | Display name of a binary operator. +binaryName :: (BinaryOp op) => op a b c -> T.Text +binaryName = opRefName binOpName . binaryId + data UnUDF a b = MkUnaryOp { unaryFn :: a -> b - , unaryName :: T.Text + , unaryId :: OpRef UnOpId , unarySymbol :: Maybe T.Text } data BinUDF a b c = MkBinaryOp { binaryFn :: a -> b -> c - , binaryName :: T.Text + , binaryId :: OpRef BinOpId , binarySymbol :: Maybe T.Text , binaryCommutative :: Bool , binaryPrecedence :: Int @@ -58,40 +70,50 @@ data BinUDF a b c = MkBinaryOp instance UnaryOp UnUDF where unaryFn (MkUnaryOp{unaryFn = f}) = f - unaryName (MkUnaryOp{unaryName = n}) = n + unaryId (MkUnaryOp{unaryId = n}) = n unarySymbol (MkUnaryOp{unarySymbol = s}) = s instance BinaryOp BinUDF where binaryFn (MkBinaryOp{binaryFn = f}) = f - binaryName (MkBinaryOp{binaryName = n}) = n + binaryId (MkBinaryOp{binaryId = n}) = n binarySymbol (MkBinaryOp{binarySymbol = s}) = s binaryCommutative (MkBinaryOp{binaryCommutative = c}) = c binaryPrecedence (MkBinaryOp{binaryPrecedence = p}) = p data AggStrategy a b where CollectAgg :: - (VG.Vector v b, Typeable v) => T.Text -> (v b -> a) -> AggStrategy a b - FoldAgg :: T.Text -> Maybe a -> (a -> b -> a) -> AggStrategy a b + (VG.Vector v b, Typeable v) => + OpRef AggId -> (v b -> a) -> AggStrategy a b + FoldAgg :: OpRef AggId -> Maybe a -> (a -> b -> a) -> AggStrategy a b MergeAgg :: (Columnable acc) => - T.Text -> + OpRef AggId -> acc -> (acc -> b -> acc) -> (acc -> acc -> acc) -> (acc -> a) -> AggStrategy a b +aggStrategyId :: AggStrategy a b -> OpRef AggId +aggStrategyId (CollectAgg i _) = i +aggStrategyId (FoldAgg i _ _) = i +aggStrategyId (MergeAgg i _ _ _ _) = i + +-- | Display name of an aggregation. +aggStrategyName :: AggStrategy a b -> T.Text +aggStrategyName = opRefName aggName . aggStrategyId + data Expr a where Col :: (Columnable a) => T.Text -> Expr a CastWith :: (Columnable a, Columnable b, Read a) => T.Text -> - T.Text -> + OpRef CastId -> (Either String a -> b) -> Expr b CastExprWith :: (Columnable a, Columnable b, Columnable src, Read a) => - T.Text -> + OpRef CastId -> (Either String a -> b) -> Expr src -> Expr b @@ -131,7 +153,7 @@ instance (Num a, Columnable a) => Num (Expr a) where Binary ( MkBinaryOp { binaryFn = (+) - , binaryName = "add" + , binaryId = Builtin BAdd , binarySymbol = Just "+" , binaryCommutative = True , binaryPrecedence = 6 @@ -143,7 +165,7 @@ instance (Num a, Columnable a) => Num (Expr a) where Binary ( MkBinaryOp { binaryFn = (-) - , binaryName = "sub" + , binaryId = Builtin BSub , binarySymbol = Just "-" , binaryCommutative = False , binaryPrecedence = 6 @@ -155,7 +177,7 @@ instance (Num a, Columnable a) => Num (Expr a) where Binary ( MkBinaryOp { binaryFn = (*) - , binaryName = "mult" + , binaryId = Builtin BMult , binarySymbol = Just "*" , binaryCommutative = True , binaryPrecedence = 7 @@ -168,15 +190,16 @@ instance (Num a, Columnable a) => Num (Expr a) where negate :: Expr a -> Expr a negate = Unary - (MkUnaryOp{unaryFn = negate, unaryName = "negate", unarySymbol = Nothing}) + (MkUnaryOp{unaryFn = negate, unaryId = Builtin UNegate, unarySymbol = Nothing}) abs :: (Num a) => Expr a -> Expr a - abs = Unary (MkUnaryOp{unaryFn = abs, unaryName = "abs", unarySymbol = Nothing}) + abs = + Unary (MkUnaryOp{unaryFn = abs, unaryId = Builtin UAbs, unarySymbol = Nothing}) signum :: (Num a) => Expr a -> Expr a signum = Unary - (MkUnaryOp{unaryFn = signum, unaryName = "signum", unarySymbol = Nothing}) + (MkUnaryOp{unaryFn = signum, unaryId = Builtin USignum, unarySymbol = Nothing}) add :: (Num a, Columnable a) => Expr a -> Expr a -> Expr a add = (+) @@ -196,7 +219,7 @@ instance (Fractional a, Columnable a) => Fractional (Expr a) where Binary ( MkBinaryOp { binaryFn = (/) - , binaryName = "divide" + , binaryId = Builtin BDivide , binarySymbol = Just "/" , binaryCommutative = False , binaryPrecedence = 7 @@ -214,80 +237,96 @@ instance (Floating a, Columnable a) => Floating (Expr a) where pi :: (Floating a, Columnable a) => Expr a pi = Lit pi exp :: (Floating a, Columnable a) => Expr a -> Expr a - exp = Unary (MkUnaryOp{unaryFn = exp, unaryName = "exp", unarySymbol = Nothing}) + exp = + Unary (MkUnaryOp{unaryFn = exp, unaryId = Builtin UExp, unarySymbol = Nothing}) sqrt :: (Floating a, Columnable a) => Expr a -> Expr a sqrt = - Unary (MkUnaryOp{unaryFn = sqrt, unaryName = "sqrt", unarySymbol = Nothing}) + Unary + (MkUnaryOp{unaryFn = sqrt, unaryId = Builtin USqrt, unarySymbol = Nothing}) (**) :: (Floating a, Columnable a) => Expr a -> Expr a -> Expr a (**) = Binary ( MkBinaryOp { binaryFn = (**) - , binaryName = "exponentiate" + , binaryId = Builtin BExponentiate , binarySymbol = Just "**" , binaryCommutative = False , binaryPrecedence = 8 } ) log :: (Floating a, Columnable a) => Expr a -> Expr a - log = Unary (MkUnaryOp{unaryFn = log, unaryName = "log", unarySymbol = Nothing}) + log = + Unary (MkUnaryOp{unaryFn = log, unaryId = Builtin ULog, unarySymbol = Nothing}) logBase :: (Floating a, Columnable a) => Expr a -> Expr a -> Expr a logBase = Binary ( MkBinaryOp { binaryFn = logBase - , binaryName = "logBase" + , binaryId = Builtin BLogBase , binarySymbol = Nothing , binaryCommutative = False , binaryPrecedence = 1 } ) sin :: (Floating a, Columnable a) => Expr a -> Expr a - sin = Unary (MkUnaryOp{unaryFn = sin, unaryName = "sin", unarySymbol = Nothing}) + sin = + Unary (MkUnaryOp{unaryFn = sin, unaryId = Builtin USin, unarySymbol = Nothing}) cos :: (Floating a, Columnable a) => Expr a -> Expr a - cos = Unary (MkUnaryOp{unaryFn = cos, unaryName = "cos", unarySymbol = Nothing}) + cos = + Unary (MkUnaryOp{unaryFn = cos, unaryId = Builtin UCos, unarySymbol = Nothing}) tan :: (Floating a, Columnable a) => Expr a -> Expr a - tan = Unary (MkUnaryOp{unaryFn = tan, unaryName = "tan", unarySymbol = Nothing}) + tan = + Unary (MkUnaryOp{unaryFn = tan, unaryId = Builtin UTan, unarySymbol = Nothing}) asin :: (Floating a, Columnable a) => Expr a -> Expr a asin = - Unary (MkUnaryOp{unaryFn = asin, unaryName = "asin", unarySymbol = Nothing}) + Unary + (MkUnaryOp{unaryFn = asin, unaryId = Builtin UAsin, unarySymbol = Nothing}) acos :: (Floating a, Columnable a) => Expr a -> Expr a acos = - Unary (MkUnaryOp{unaryFn = acos, unaryName = "acos", unarySymbol = Nothing}) + Unary + (MkUnaryOp{unaryFn = acos, unaryId = Builtin UAcos, unarySymbol = Nothing}) atan :: (Floating a, Columnable a) => Expr a -> Expr a atan = - Unary (MkUnaryOp{unaryFn = atan, unaryName = "atan", unarySymbol = Nothing}) + Unary + (MkUnaryOp{unaryFn = atan, unaryId = Builtin UAtan, unarySymbol = Nothing}) sinh :: (Floating a, Columnable a) => Expr a -> Expr a sinh = - Unary (MkUnaryOp{unaryFn = sinh, unaryName = "sinh", unarySymbol = Nothing}) + Unary + (MkUnaryOp{unaryFn = sinh, unaryId = Builtin USinh, unarySymbol = Nothing}) cosh :: (Floating a, Columnable a) => Expr a -> Expr a cosh = - Unary (MkUnaryOp{unaryFn = cosh, unaryName = "cosh", unarySymbol = Nothing}) + Unary + (MkUnaryOp{unaryFn = cosh, unaryId = Builtin UCosh, unarySymbol = Nothing}) asinh :: (Floating a, Columnable a) => Expr a -> Expr a asinh = Unary - (MkUnaryOp{unaryFn = asinh, unaryName = "asinh", unarySymbol = Nothing}) + (MkUnaryOp{unaryFn = asinh, unaryId = Builtin UAsinh, unarySymbol = Nothing}) acosh :: (Floating a, Columnable a) => Expr a -> Expr a acosh = Unary - (MkUnaryOp{unaryFn = acosh, unaryName = "acosh", unarySymbol = Nothing}) + (MkUnaryOp{unaryFn = acosh, unaryId = Builtin UAcosh, unarySymbol = Nothing}) atanh :: (Floating a, Columnable a) => Expr a -> Expr a atanh = Unary - (MkUnaryOp{unaryFn = atanh, unaryName = "atanh", unarySymbol = Nothing}) + (MkUnaryOp{unaryFn = atanh, unaryId = Builtin UAtanh, unarySymbol = Nothing}) instance (Show a) => Show (Expr a) where show :: Expr a -> String show (Col name) = "(col @" ++ show (typeRep @a) ++ " " ++ show name ++ ")" - show (CastWith name tag _) = "(castWith " ++ show tag ++ " " ++ show name ++ ")" - show (CastExprWith tag _ inner) = "(castExprWith " ++ show tag ++ " " ++ show inner ++ ")" + show (CastWith name tag _) = + "(castWith " ++ show (opRefName castName tag) ++ " " ++ show name ++ ")" + show (CastExprWith tag _ inner) = + "(castExprWith " + ++ show (opRefName castExprName tag) + ++ " " + ++ show inner + ++ ")" show (Lit value) = "(lit (" ++ show value ++ "))" show (If cond l r) = "(ifThenElse " ++ show cond ++ " " ++ show l ++ " " ++ show r ++ ")" show (Unary op value) = "(" ++ T.unpack (unaryName op) ++ " " ++ show value ++ ")" show (Binary op a b) = "(" ++ T.unpack (binaryName op) ++ " " ++ show a ++ " " ++ show b ++ ")" - show (Agg (CollectAgg op _) expr) = "(" ++ T.unpack op ++ " " ++ show expr ++ ")" - show (Agg (FoldAgg op _ _) expr) = "(" ++ T.unpack op ++ " " ++ show expr ++ ")" - show (Agg (MergeAgg op _ _ _ _) expr) = "(" ++ T.unpack op ++ " " ++ show expr ++ ")" + show (Agg strat expr) = + "(" ++ T.unpack (aggStrategyName strat) ++ " " ++ show expr ++ ")" show (Over keys inner) = "(over " ++ show keys ++ " " ++ show inner ++ ")" normalize :: (Show a, Typeable a) => Expr a -> Expr a @@ -318,41 +357,77 @@ compareExpr e1 e2 = compare (exprKey e1) (exprKey e2) where exprKey :: Expr a -> String exprKey (Col name) = "0:" ++ T.unpack name - exprKey (CastWith name tag _) = "0CW:" ++ T.unpack name ++ ":" ++ T.unpack tag - exprKey (CastExprWith tag _ _) = "0CE:" ++ T.unpack tag + exprKey (CastWith name tag _) = "0CW:" ++ T.unpack name ++ ":" ++ idKey tag + exprKey (CastExprWith tag _ _) = "0CE:" ++ idKey tag exprKey (Lit val) = "1:" ++ show val exprKey (If c t e) = "2:" ++ exprKey c ++ exprKey t ++ exprKey e - exprKey (Unary op e) = "3:" ++ T.unpack (unaryName op) ++ exprKey e - exprKey (Binary op e1' e2') = "4:" ++ T.unpack (binaryName op) ++ exprKey e1' ++ exprKey e2' - exprKey (Agg (CollectAgg name _) e) = "5:" ++ T.unpack name ++ exprKey e - exprKey (Agg (FoldAgg name _ _) e) = "5:" ++ T.unpack name ++ exprKey e - exprKey (Agg (MergeAgg name _ _ _ _) e) = "5:" ++ T.unpack name ++ exprKey e + exprKey (Unary op e) = "3:" ++ idKey (unaryId op) ++ exprKey e + exprKey (Binary op e1' e2') = + "4:" ++ idKey (binaryId op) ++ exprKey e1' ++ exprKey e2' + exprKey (Agg strat e) = "5:" ++ idKey (aggStrategyId strat) ++ exprKey e exprKey (Over keys e) = "6:over:" ++ show keys ++ exprKey e - + idKey :: (Show k) => OpRef k -> String + idKey (Builtin k) = "b:" ++ show k + idKey (Named n) = "n:" ++ T.unpack n + idKey (Anonymous l) = "a:" ++ T.unpack l + +-- | A key equal only for provably equal expressions; Nothing if anonymous. +structuralKey :: (Columnable a) => Expr a -> Maybe String +structuralKey = go . normalize + where + go :: forall x. Expr x -> Maybe String + go e = case e of + Col name -> Just ("c" ++ show name ++ "@" ++ show (typeRep @x)) + CastWith name tag _ -> + idKey tag >>= \k -> Just ("w" ++ show name ++ k ++ "@" ++ show (typeRep @x)) + CastExprWith tag _ inner -> idKey tag >>= \k -> wrap ("x" ++ k) [go inner] + Lit v -> Just ("l" ++ show v ++ "@" ++ show (typeRep @x)) + If c t f -> wrap "if" [go c, go t, go f] + Unary op inner -> idKey (unaryId op) >>= \k -> wrap k [go inner] + Binary op l r -> idKey (binaryId op) >>= \k -> wrap k [go l, go r] + Agg strat inner -> + idKey (aggStrategyId strat) >>= \k -> wrap ("agg" ++ k) [go inner] + Over keys inner -> wrap ("over" ++ show keys) [go inner] + wrap tag parts = (\ps -> tag ++ "(" ++ concatMap (++ ",") ps ++ ")") <$> sequence parts + idKey :: (Show k) => OpRef k -> Maybe String + idKey (Builtin k) = Just ("b" ++ show k) + idKey (Named n) = Just ("n" ++ show n) + idKey (Anonymous _) = Nothing + +-- | Provable equality; anonymous operators are never equal. eqExpr :: forall a. (Columnable a) => Expr a -> Expr a -> Bool eqExpr l r = eqNormalized (normalize l) (normalize r) + +eqNormalized :: forall a. Expr a -> Expr a -> Bool +eqNormalized = go where - exprEq :: (Columnable b, Columnable c) => Expr b -> Expr c -> Bool - exprEq e1 e2 = case testEquality (typeOf e1) (typeOf e2) of - Just Refl -> eqExpr e1 e2 + sub :: forall b c. (Columnable b, Columnable c) => Expr b -> Expr c -> Bool + sub e1 e2 = case testEquality (typeRep @b) (typeRep @c) of + Just Refl -> go e1 e2 Nothing -> False - eqNormalized :: Expr a -> Expr a -> Bool - eqNormalized (Col n1) (Col n2) = n1 == n2 - eqNormalized (CastWith n1 t1 _) (CastWith n2 t2 _) = n1 == n2 && t1 == t2 - eqNormalized (CastExprWith t1 _ e1) (CastExprWith t2 _ e2) = t1 == t2 && e1 `exprEq` e2 - eqNormalized (Lit v1) (Lit v2) = v1 == v2 - eqNormalized (If c1 t1 e1) (If c2 t2 e2) = - eqExpr c1 c2 && t1 `exprEq` t2 && e1 `exprEq` e2 - eqNormalized (Unary op1 e1) (Unary op2 e2) = unaryName op1 == unaryName op2 && e1 `exprEq` e2 - eqNormalized (Binary op1 e1a e1b) (Binary op2 e2a e2b) = binaryName op1 == binaryName op2 && e1a `exprEq` e2a && e1b `exprEq` e2b - eqNormalized (Agg (CollectAgg n1 _) e1) (Agg (CollectAgg n2 _) e2) = - n1 == n2 && e1 `exprEq` e2 - eqNormalized (Agg (FoldAgg n1 _ _) e1) (Agg (FoldAgg n2 _ _) e2) = - n1 == n2 && e1 `exprEq` e2 - eqNormalized (Agg (MergeAgg n1 _ _ _ _) e1) (Agg (MergeAgg n2 _ _ _ _) e2) = - n1 == n2 && e1 `exprEq` e2 - eqNormalized (Over k1 e1) (Over k2 e2) = k1 == k2 && e1 `exprEq` e2 - eqNormalized _ _ = False + go :: forall x. Expr x -> Expr x -> Bool + go (Col n1) (Col n2) = n1 == n2 + go (CastWith n1 t1 _) (CastWith n2 t2 _) = n1 == n2 && sameOp t1 t2 + go (CastExprWith t1 _ e1) (CastExprWith t2 _ e2) = sameOp t1 t2 && e1 `sub` e2 + go (Lit v1) (Lit v2) = v1 == v2 + go (If c1 t1 e1) (If c2 t2 e2) = go c1 c2 && go t1 t2 && go e1 e2 + go (Unary op1 e1) (Unary op2 e2) = + sameOp (unaryId op1) (unaryId op2) && e1 `sub` e2 + go (Binary op1 e1a e1b) (Binary op2 e2a e2b) = + sameOp (binaryId op1) (binaryId op2) + && e1a `sub` e2a + && e1b `sub` e2b + go (Agg s1 e1) (Agg s2 e2) = + sameStrategy s1 s2 + && sameOp (aggStrategyId s1) (aggStrategyId s2) + && e1 `sub` e2 + go (Over k1 e1) (Over k2 e2) = k1 == k2 && go e1 e2 + go _ _ = False + sameStrategy :: AggStrategy p q -> AggStrategy p' q' -> Bool + sameStrategy CollectAgg{} CollectAgg{} = True + sameStrategy FoldAgg{} FoldAgg{} = True + sameStrategy MergeAgg{} MergeAgg{} = True + sameStrategy _ _ = False replaceExpr :: forall a b. @@ -435,7 +510,8 @@ prettyPrintWidth width = P.render width . toDoc 0 toDoc prec expr = case expr of Col name -> P.text (T.unpack name) CastWith name _ _ -> P.text (T.unpack name) - CastExprWith tag _ inner -> P.text (T.unpack tag) <> P.parens (toDoc 0 inner) + CastExprWith tag _ inner -> + P.text (T.unpack (opRefName castExprName tag)) <> P.parens (toDoc 0 inner) Lit value -> P.text (show value) If{} -> renderIf prec expr Unary op arg -> @@ -446,9 +522,8 @@ prettyPrintWidth width = P.render width . toDoc 0 Nothing -> P.text (T.unpack (binaryName op)) <> P.parens (toDoc 0 l <> P.text ", " <> toDoc 0 r) - Agg (CollectAgg op _) arg -> P.text (T.unpack op) <> P.parens (toDoc 0 arg) - Agg (FoldAgg op _ _) arg -> P.text (T.unpack op) <> P.parens (toDoc 0 arg) - Agg (MergeAgg op _ _ _ _) arg -> P.text (T.unpack op) <> P.parens (toDoc 0 arg) + Agg strat arg -> + P.text (T.unpack (aggStrategyName strat)) <> P.parens (toDoc 0 arg) Over keys inner -> toDoc 0 inner <> P.text (".over(" ++ show (map T.unpack keys) ++ ")") diff --git a/dataframe-core/src-internal/DataFrame/Internal/Expression/OpId.hs b/dataframe-core/src-internal/DataFrame/Internal/Expression/OpId.hs new file mode 100644 index 00000000..56658ab2 --- /dev/null +++ b/dataframe-core/src-internal/DataFrame/Internal/Expression/OpId.hs @@ -0,0 +1,260 @@ +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE OverloadedStrings #-} + +-- | What an operator computes; kernels key on these tags, not display names. +module DataFrame.Internal.Expression.OpId ( + OpRef (..), + UnOpId (..), + BinOpId (..), + AggId (..), + CastId (..), + sameOp, + builtinOf, + opRefName, + unOpName, + binOpName, + aggName, + castName, + castExprName, +) where + +import qualified Data.Text as T + +-- | Builtin: library op; Named: equal names, equal functions; Anonymous: label. +data OpRef k + = Builtin !k + | Named !T.Text + | Anonymous !T.Text + deriving (Show, Functor) + +-- | Provable equality: 'Anonymous' is never equal to anything, even itself. +sameOp :: (Eq k) => OpRef k -> OpRef k -> Bool +sameOp (Builtin a) (Builtin b) = a == b +sameOp (Named a) (Named b) = a == b +sameOp _ _ = False +{-# INLINE sameOp #-} + +-- | The builtin tag, if any; the only thing kernels may dispatch on. +builtinOf :: OpRef k -> Maybe k +builtinOf (Builtin k) = Just k +builtinOf _ = Nothing +{-# INLINE builtinOf #-} + +-- | The display name: the builtin's name, the given name, or the label. +opRefName :: (k -> T.Text) -> OpRef k -> T.Text +opRefName f (Builtin k) = f k +opRefName _ (Named n) = n +opRefName _ (Anonymous l) = l + +-- | Library unary operators. Parameters are part of the identity. +data UnOpId + = UNegate + | UAbs + | USignum + | UExp + | USqrt + | ULog + | USin + | UCos + | UTan + | UAsin + | UAcos + | UAtan + | USinh + | UCosh + | UAsinh + | UAcosh + | UAtanh + | UNot + | UToDouble + | URelu + | UToMaybe + | -- | The default, rendered with 'show'. + UFromMaybe !T.Text + | UIsJust + | UIsNothing + | UFromJust + | UFirstOrNothing + | ULastOrNothing + | UTextLength + | USplitOn !T.Text + | UMatch !T.Text + | UMatchAll !T.Text + | UParseDate !T.Text + | -- | The mapping, rendered with 'show'. + URecode !T.Text + | -- | The default and the mapping, rendered with 'show'. + URecodeWithDefault !T.Text !T.Text + deriving (Eq, Ord, Show) + +-- | Library binary operators. +data BinOpId + = BAdd + | BSub + | BMult + | BDivide + | BExponentiate + | BLogBase + | BDiv + | BMod + | -- | @(^)@: integral exponent. + BPow + | -- | @(^^)@: fractional base, integral exponent. + BPowFrac + | BEq + | BNeq + | BLt + | BLeq + | BGt + | BGeq + | BAnd + | BOr + | BNullAnd + | BNullOr + | BNullAdd + | BNullSub + | BNullMul + | BNullDiv + | BMin + | BMax + | BContains + | BDaysBetween + deriving (Eq, Ord, Show, Enum, Bounded) + +-- | Library aggregations. +data AggId + = ASum + | AMinimum + | AMaximum + | AMean + | ACount + | ACountDistinct + | AMode + | ACollect + | ASumMaybe + | AMeanMaybe + | AMedianMaybe + | AVariance + | AMedian + | AStddev + | AStddevMaybe + | APercentile !Int + | ATopK !Int + | ABottomK !Int + deriving (Eq, Ord, Show) + +-- | Display and wire name of a unary builtin. +unOpName :: UnOpId -> T.Text +unOpName op = case op of + UNegate -> "negate" + UAbs -> "abs" + USignum -> "signum" + UExp -> "exp" + USqrt -> "sqrt" + ULog -> "log" + USin -> "sin" + UCos -> "cos" + UTan -> "tan" + UAsin -> "asin" + UAcos -> "acos" + UAtan -> "atan" + USinh -> "sinh" + UCosh -> "cosh" + UAsinh -> "asinh" + UAcosh -> "acosh" + UAtanh -> "atanh" + UNot -> "not" + UToDouble -> "toDouble" + URelu -> "relu" + UToMaybe -> "toMaybe" + UFromMaybe _ -> "fromMaybe" + UIsJust -> "isJust" + UIsNothing -> "isNothing" + UFromJust -> "fromJust" + UFirstOrNothing -> "firstOrNothing" + ULastOrNothing -> "lastOrNothing" + UTextLength -> "textLength" + USplitOn _ -> "splitOn" + UMatch r -> "match " <> T.pack (show r) + UMatchAll r -> "matchAll " <> T.pack (show r) + UParseDate f -> "parseDate " <> f + URecode m -> "recode " <> m + URecodeWithDefault d m -> "recodeWithDefault " <> d <> " " <> m + +-- | Display and wire name of a binary builtin. +binOpName :: BinOpId -> T.Text +binOpName op = case op of + BAdd -> "add" + BSub -> "sub" + BMult -> "mult" + BDivide -> "divide" + BExponentiate -> "exponentiate" + BLogBase -> "logBase" + BDiv -> "div" + BMod -> "mod" + BPow -> "pow" + BPowFrac -> "pow" + BEq -> "eq" + BNeq -> "neq" + BLt -> "lt" + BLeq -> "leq" + BGt -> "gt" + BGeq -> "geq" + BAnd -> "and" + BOr -> "or" + BNullAnd -> "nulland" + BNullOr -> "nullor" + BNullAdd -> "nulladd" + BNullSub -> "nullsub" + BNullMul -> "nullmul" + BNullDiv -> "nulldiv" + BMin -> "min" + BMax -> "max" + BContains -> "contains" + BDaysBetween -> "daysBetween" + +-- | Display and wire name of an aggregation builtin. +aggName :: AggId -> T.Text +aggName op = case op of + ASum -> "sum" + AMinimum -> "minimum" + AMaximum -> "maximum" + AMean -> "mean" + ACount -> "count" + ACountDistinct -> "countDistinct" + AMode -> "mode" + ACollect -> "collect" + ASumMaybe -> "sumMaybe" + AMeanMaybe -> "meanMaybe" + AMedianMaybe -> "medianMaybe" + AVariance -> "variance" + AMedian -> "median" + AStddev -> "stddev" + AStddevMaybe -> "stddevMaybe" + APercentile n -> "percentile " <> T.pack (show n) + ATopK k -> "topK_" <> T.pack (show k) + ABottomK k -> "bottomK_" <> T.pack (show k) + +-- | Library casts; a default value is shown so different defaults differ. +data CastId + = CCast + | CCastWithDefault !T.Text + | CCastEither + | CUnsafeCast + deriving (Eq, Ord, Show) + +-- | Display name of a cast reading a column by name. +castName :: CastId -> T.Text +castName c = case c of + CCast -> "cast" + CCastWithDefault d -> "castWithDefault:" <> d + CCastEither -> "castEither" + CUnsafeCast -> "unsafeCast" + +-- | Display name of a cast applied to an expression. +castExprName :: CastId -> T.Text +castExprName c = case c of + CCast -> "castExpr" + CCastWithDefault d -> "castExprWithDefault:" <> d + CCastEither -> "castExprEither" + CUnsafeCast -> "unsafeCastExpr" diff --git a/dataframe-core/src-internal/DataFrame/Internal/Expression/Operators.hs b/dataframe-core/src-internal/DataFrame/Internal/Expression/Operators.hs index e79fa9ed..1d141c0d 100644 --- a/dataframe-core/src-internal/DataFrame/Internal/Expression/Operators.hs +++ b/dataframe-core/src-internal/DataFrame/Internal/Expression/Operators.hs @@ -11,17 +11,20 @@ import qualified Data.Text as T import DataFrame.Internal.Column (Columnable) import DataFrame.Internal.Column.Types (Promote, PromoteDiv) import DataFrame.Internal.Expression ( + BinOpId (..), BinUDF (MkBinaryOp), BinaryOp ( binaryCommutative, binaryFn, - binaryName, + binaryId, binaryPrecedence, binarySymbol ), Expr (Binary, Col, If, Lit, Unary), NamedExpr, + OpRef (..), SomeExpr (SomeExpr), + UnOpId, UnUDF (MkUnaryOp), ) import DataFrame.Internal.Expression.Operators.Nullable ( @@ -77,10 +80,16 @@ lit = Lit (.=) :: (Columnable a) => T.Text -> Expr a -> NamedExpr (.=) = flip as +-- | A user op; equal names must mean equal functions. Never hits builtins. liftDecorated :: (Columnable a, Columnable b) => (a -> b) -> T.Text -> Maybe T.Text -> Expr a -> Expr b -liftDecorated f opName rep = Unary (MkUnaryOp f opName rep) +liftDecorated f opName rep = Unary (MkUnaryOp f (Named opName) rep) + +liftBuiltin :: + (Columnable a, Columnable b) => + UnOpId -> (a -> b) -> Maybe T.Text -> Expr a -> Expr b +liftBuiltin op f rep = Unary (MkUnaryOp f (Builtin op) rep) lift2Decorated :: (Columnable c, Columnable b, Columnable a) => @@ -93,7 +102,19 @@ lift2Decorated :: Expr b -> Expr a lift2Decorated f opName rep comm prec = - Binary (MkBinaryOp f opName rep comm prec) + Binary (MkBinaryOp f (Named opName) rep comm prec) + +lift2Builtin :: + (Columnable c, Columnable b, Columnable a) => + BinOpId -> + (c -> b -> a) -> + Maybe T.Text -> + Bool -> + Int -> + Expr c -> + Expr b -> + Expr a +lift2Builtin op f rep comm prec = Binary (MkBinaryOp f (Builtin op) rep comm prec) data NullEq a b c where NullEq :: @@ -155,45 +176,45 @@ data NullOr a b c where instance BinaryOp NullEq where binaryFn NullEq = applyNull2 (widenCmpOp (==)) - binaryName NullEq = "eq" + binaryId NullEq = Builtin BEq binarySymbol NullEq = Just ".==" binaryCommutative NullEq = True binaryPrecedence NullEq = 4 instance BinaryOp NullNeq where binaryFn NullNeq = applyNull2 (widenCmpOp (/=)) - binaryName NullNeq = "neq" + binaryId NullNeq = Builtin BNeq binarySymbol NullNeq = Just "./=" binaryCommutative NullNeq = True binaryPrecedence NullNeq = 4 instance BinaryOp NullLt where binaryFn NullLt = applyNull2 (widenCmpOp (<)) - binaryName NullLt = "lt" + binaryId NullLt = Builtin BLt binarySymbol NullLt = Just ".<" binaryPrecedence NullLt = 4 instance BinaryOp NullGt where binaryFn NullGt = applyNull2 (widenCmpOp (>)) - binaryName NullGt = "gt" + binaryId NullGt = Builtin BGt binarySymbol NullGt = Just ".>" binaryPrecedence NullGt = 4 instance BinaryOp NullLeq where binaryFn NullLeq = applyNull2 (widenCmpOp (<=)) - binaryName NullLeq = "leq" + binaryId NullLeq = Builtin BLeq binarySymbol NullLeq = Just ".<=" binaryPrecedence NullLeq = 4 instance BinaryOp NullGeq where binaryFn NullGeq = applyNull2 (widenCmpOp (>=)) - binaryName NullGeq = "geq" + binaryId NullGeq = Builtin BGeq binarySymbol NullGeq = Just ".>=" binaryPrecedence NullGeq = 4 instance BinaryOp NullAnd where binaryFn NullAnd = nullCmpOp (&&) - binaryName NullAnd = "nulland" + binaryId NullAnd = Builtin BNullAnd binarySymbol NullAnd = Just ".&&" binaryCommutative NullAnd = True binaryPrecedence NullAnd = 3 instance BinaryOp NullOr where binaryFn NullOr = nullCmpOp (||) - binaryName NullOr = "nullor" + binaryId NullOr = Builtin BNullOr binarySymbol NullOr = Just ".||" binaryCommutative NullOr = True binaryPrecedence NullOr = 2 @@ -203,42 +224,42 @@ instance BinaryOp NullOr where Expr a -> Expr a -> Expr Bool -(.==.) = lift2Decorated (==) "eq" (Just ".==.") True 4 +(.==.) = lift2Builtin BEq (==) (Just ".==.") True 4 (./=.) :: (Columnable a, Eq a) => Expr a -> Expr a -> Expr Bool -(./=.) = lift2Decorated (/=) "neq" (Just "./=.") True 4 +(./=.) = lift2Builtin BNeq (/=) (Just "./=.") True 4 (.<.) :: (Columnable a, Ord a) => Expr a -> Expr a -> Expr Bool -(.<.) = lift2Decorated (<) "lt" (Just ".<.") False 4 +(.<.) = lift2Builtin BLt (<) (Just ".<.") False 4 (.>.) :: (Columnable a, Ord a) => Expr a -> Expr a -> Expr Bool -(.>.) = lift2Decorated (>) "gt" (Just ".>.") False 4 +(.>.) = lift2Builtin BGt (>) (Just ".>.") False 4 (.<=.) :: (Columnable a, Ord a) => Expr a -> Expr a -> Expr Bool -(.<=.) = lift2Decorated (<=) "leq" (Just ".<=.") False 4 +(.<=.) = lift2Builtin BLeq (<=) (Just ".<=.") False 4 (.>=.) :: (Columnable a, Ord a) => Expr a -> Expr a -> Expr Bool -(.>=.) = lift2Decorated (>=) "geq" (Just ".>=.") False 4 +(.>=.) = lift2Builtin BGeq (>=) (Just ".>=.") False 4 (.+.) :: (Columnable a, Num a) => Expr a -> Expr a -> Expr a (.+.) = (+) @@ -263,7 +284,7 @@ instance BinaryOp NullOr where Expr a -> Expr b -> Expr (WidenResult a b) -(.+) = lift2Decorated (applyNull2 (widenArithOp (+))) "nulladd" (Just ".+") True 6 +(.+) = lift2Builtin BNullAdd (applyNull2 (widenArithOp (+))) (Just ".+") True 6 -- | Nullable-aware subtraction. (.-) :: @@ -274,7 +295,7 @@ instance BinaryOp NullOr where Expr a -> Expr b -> Expr (WidenResult a b) -(.-) = lift2Decorated (applyNull2 (widenArithOp (-))) "nullsub" (Just ".-") False 6 +(.-) = lift2Builtin BNullSub (applyNull2 (widenArithOp (-))) (Just ".-") False 6 -- | Nullable-aware multiplication. (.*) :: @@ -285,7 +306,7 @@ instance BinaryOp NullOr where Expr a -> Expr b -> Expr (WidenResult a b) -(.*) = lift2Decorated (applyNull2 (widenArithOp (*))) "nullmul" (Just ".*") True 7 +(.*) = lift2Builtin BNullMul (applyNull2 (widenArithOp (*))) (Just ".*") True 7 -- | Nullable-aware division. Integral operands are promoted to Double. (./) :: @@ -296,7 +317,7 @@ instance BinaryOp NullOr where Expr a -> Expr b -> Expr (WidenResultDiv a b) -(./) = lift2Decorated (applyNull2 (divArithOp (/))) "nulldiv" (Just "./") False 7 +(./) = lift2Builtin BNullDiv (applyNull2 (divArithOp (/))) (Just "./") False 7 -- Three-valued comparisons: Nothing if either operand is Nothing. @@ -367,10 +388,10 @@ instance BinaryOp NullOr where (.>=) = Binary NullGeq (.&&.) :: Expr Bool -> Expr Bool -> Expr Bool -(.&&.) = lift2Decorated (&&) "and" (Just ".&&.") True 3 +(.&&.) = lift2Builtin BAnd (&&) (Just ".&&.") True 3 (.||.) :: Expr Bool -> Expr Bool -> Expr Bool -(.||.) = lift2Decorated (||) "or" (Just ".||.") True 2 +(.||.) = lift2Builtin BOr (||) (Just ".||.") True 2 -- | Nullable-aware logical AND, giving @Maybe Bool@ if either side is nullable. (.&&) :: @@ -398,7 +419,7 @@ instance BinaryOp NullOr where , Num (Promote (BaseType a) (BaseType b)) ) => Expr a -> Expr b -> Expr a -(.^^) = lift2Decorated (applyNull2 (^^)) "pow" (Just ".^^") False 8 +(.^^) = lift2Builtin BPowFrac (applyNull2 (^^)) (Just ".^^") False 8 (.^) :: ( Columnable (BaseType a) @@ -410,16 +431,16 @@ instance BinaryOp NullOr where , Num (Promote (BaseType a) (BaseType b)) ) => Expr a -> Expr b -> Expr a -(.^) = lift2Decorated (applyNull2 (^)) "pow" (Just ".^") False 8 +(.^) = lift2Builtin BPow (applyNull2 (^)) (Just ".^") False 8 -- Same-type (non-nullable) exponentiation operators (.^^.) :: (Columnable a, Columnable b, Fractional a, Integral b) => Expr a -> Expr b -> Expr a -(.^^.) = lift2Decorated (^^) "pow" (Just ".^^.") False 8 +(.^^.) = lift2Builtin BPowFrac (^^) (Just ".^^.") False 8 (.^.) :: (Columnable a, Columnable b, Num a, Integral b) => Expr a -> Expr b -> Expr a -(.^.) = lift2Decorated (^) "pow" (Just ".^.") False 8 +(.^.) = lift2Builtin BPow (^) (Just ".^.") False 8 diff --git a/dataframe-core/src-internal/DataFrame/Internal/Expression/Simplify.hs b/dataframe-core/src-internal/DataFrame/Internal/Expression/Simplify.hs index 21e5d52d..c5668338 100644 --- a/dataframe-core/src-internal/DataFrame/Internal/Expression/Simplify.hs +++ b/dataframe-core/src-internal/DataFrame/Internal/Expression/Simplify.hs @@ -27,9 +27,11 @@ import DataFrame.Internal.Column (Columnable) import DataFrame.Internal.Expression ( BinaryOp, Expr (..), - UnaryOp (unaryName), + OpRef (..), + UnOpId (..), + UnaryOp (unaryId), eqExpr, - normalize, + structuralKey, ) import DataFrame.Internal.Expression.Operators ( NullAnd, @@ -80,7 +82,7 @@ simplifyB expr = case expr of Unary (op :: op b a) inner | Just Refl <- testEquality (typeRep @a) (typeRep @Bool) , Just Refl <- testEquality (typeRep @b) (typeRep @Bool) - , unaryName op == "not" -> + , Builtin UNot <- unaryId op -> simplifyNot op (simplifyB inner) | otherwise -> expr If c t f -> @@ -104,7 +106,7 @@ simplifyNot op inner = case asBoolLit inner of Just b -> Lit (not b) Nothing -> case inner of Unary (op2 :: op2 b2 Bool) inner2 - | unaryName op2 == "not" + | Builtin UNot <- unaryId op2 , Just Refl <- testEquality (typeRep @b2) (typeRep @Bool) -> inner2 _ -> Unary op inner @@ -214,7 +216,7 @@ litDouble _ = Nothing -- | True for a column lifted from an integral type (never NaN). integralColE :: forall c. (Columnable c) => Expr c -> Bool -integralColE (Unary op _) = unaryName op == "toDouble" +integralColE (Unary op _) | Builtin UToDouble <- unaryId op = True integralColE _ = or [ matches @Int @@ -226,15 +228,17 @@ integralColE _ = atomOf :: forall a. (Columnable a) => Expr a -> Maybe Atom atomOf (Unary fm (Binary (op :: op c b r) (colE :: Expr c) litE)) - | unaryName fm == "fromMaybe" + | Builtin (UFromMaybe "False") <- unaryId fm , Just cmp <- cmpOf op - , Just t <- litDouble litE = - Just (Atom cmp t (show (normalize colE)) FalseOnNull (integralColE colE)) + , Just t <- litDouble litE + , Just key <- structuralKey colE = + Just (Atom cmp t key FalseOnNull (integralColE colE)) atomOf (Binary (op :: op c b a) (colE :: Expr c) litE) | Just cmp <- cmpOf op - , Just t <- litDouble litE = + , Just t <- litDouble litE + , Just key <- structuralKey colE = let nk = if isMaybeTy @c then UnknownOnNull else Total - in Just (Atom cmp t (show (normalize colE)) nk (integralColE colE)) + in Just (Atom cmp t key nk (integralColE colE)) atomOf _ = Nothing simplifyPredicatePair :: diff --git a/dataframe-core/src-internal/DataFrame/Internal/Interpreter.hs b/dataframe-core/src-internal/DataFrame/Internal/Interpreter.hs index 7e54f0b9..a722862c 100644 --- a/dataframe-core/src-internal/DataFrame/Internal/Interpreter.hs +++ b/dataframe-core/src-internal/DataFrame/Internal/Interpreter.hs @@ -383,7 +383,7 @@ fromOperand (K.Column v) = Flat (UnboxedColumn Nothing v) builtinBinary :: forall c b a. (Columnable c, Columnable b, Columnable a) => - T.Text -> Maybe (Value c -> Value b -> Maybe (Value a)) + BinOpId -> Maybe (Value c -> Value b -> Maybe (Value a)) builtinBinary name = do Refl <- testEquality (typeRep @c) (typeRep @b) case (comparison, arithmetic, integral, textEquality, textContains) of @@ -399,8 +399,8 @@ builtinBinary name = do Refl <- testEquality (typeRep @a) (typeRep @Int) Refl <- testEquality (typeRep @c) (typeRep @Int) op <- case name of - "div" -> Just div - "mod" -> Just mod + BDiv -> Just div + BMod -> Just mod _ -> Nothing -- Only non-zero literal divisors: lazy evaluation might skip a zero. Just $ \l r -> case r of @@ -411,18 +411,18 @@ builtinBinary name = do Refl <- testEquality (typeRep @c) (typeRep @T.Text) Refl <- testEquality (typeRep @a) (typeRep @Bool) case name of - "contains" -> Just packedContains + BContains -> Just packedContains _ -> Nothing comparison :: (c ~ b) => Maybe (Value c -> Value c -> Maybe (Value a)) comparison = do Refl <- testEquality (typeRep @a) (typeRep @Bool) - op <- K.cmpOpByName name + op <- K.cmpOpOf name run <- K.compareOperands @c Just (\l r -> fromOperand <$> (run op <$> operand l <*> operand r)) arithmetic :: (c ~ b) => Maybe (Value c -> Value c -> Maybe (Value a)) arithmetic = do Refl <- testEquality (typeRep @a) (typeRep @c) - op <- K.arithOpByName name + op <- K.arithOpOf name run <- K.arithOperands @c case sUnbox @c of STrue -> Just (\l r -> fromOperand <$> (run op <$> operand l <*> operand r)) @@ -432,8 +432,8 @@ builtinBinary name = do Refl <- testEquality (typeRep @c) (typeRep @T.Text) Refl <- testEquality (typeRep @a) (typeRep @Bool) wantEq <- case name of - "eq" -> Just True - "neq" -> Just False + BEq -> Just True + BNeq -> Just False _ -> Nothing Just (packedEquals wantEq) @@ -458,14 +458,14 @@ packedEquals wantEq l r = do -- | As 'builtinBinary', for unary operators. builtinUnary :: forall b a. - (Columnable b, Columnable a) => T.Text -> Maybe (Value b -> Maybe (Value a)) -builtinUnary "toDouble" + (Columnable b, Columnable a) => UnOpId -> Maybe (Value b -> Maybe (Value a)) +builtinUnary UToDouble | Just Refl <- testEquality (typeRep @b) (typeRep @Int) , Just Refl <- testEquality (typeRep @a) (typeRep @Double) = Just $ \v -> case operand v of Just (K.Column x) -> Just (Flat (UnboxedColumn Nothing (K.intToDouble x))) _ -> Nothing -builtinUnary "textLength" +builtinUnary UTextLength | Just Refl <- testEquality (typeRep @b) (typeRep @T.Text) , Just Refl <- testEquality (typeRep @a) (typeRep @Int) = Just $ \case @@ -477,15 +477,15 @@ applyUnary :: forall op b a. (UnaryOp op, Columnable b, Columnable a) => op b a -> Value b -> Either DataFrameException (Value a) -applyUnary op v = case builtinUnary @b @a (unaryName op) >>= ($ v) of +applyUnary op v = case builtinOf (unaryId op) >>= builtinUnary @b @a >>= ($ v) of Just out -> Right out - Nothing -> liftValue (fastUnaryFn @b @a (unaryName op) (unaryFn op)) v + Nothing -> liftValue (fastUnaryFn @b @a (unaryId op) (unaryFn op)) v applyBinary :: forall op c b a. (BinaryOp op, Columnable c, Columnable b, Columnable a) => op c b a -> Value c -> Value b -> Either DataFrameException (Value a) -applyBinary op l r = case builtinBinary @c @b @a (binaryName op) >>= (\run -> run l r) of +applyBinary op l r = case builtinOf (binaryId op) >>= builtinBinary @c @b @a >>= (\run -> run l r) of Just out -> Right out Nothing -> liftValue2 (binaryFn op) l r @@ -519,8 +519,8 @@ narrow env positions = env{envRows = Rows rows} data Key = KeyColumn T.Text SomeTypeRep | KeyLiteral Literal - | KeyUnary T.Text SomeTypeRep Key - | KeyBinary T.Text SomeTypeRep Key Key + | KeyUnary UnOpId SomeTypeRep Key + | KeyBinary BinOpId SomeTypeRep Key Key deriving (Eq, Ord) data Literal @@ -544,13 +544,15 @@ keyOf expr = fst <$> go maxKeyNodes expr Col name -> Just (KeyColumn name (SomeTypeRep (typeRep @x)), budget - 1) Lit v -> (\l -> (KeyLiteral l, budget - 1)) <$> literalKey v Unary op (inner :: Expr b) -> do - _ <- builtinUnary @b @x (unaryName op) + tag <- builtinOf (unaryId op) + _ <- builtinUnary @b @x tag (k, left) <- go (budget - 1) inner - Just (KeyUnary (unaryName op) (SomeTypeRep (typeRep @b)) k, left) + Just (KeyUnary tag (SomeTypeRep (typeRep @b)) k, left) Binary op (l :: Expr c) (r :: Expr b) -> do - _ <- builtinBinary @c @b @x (binaryName op) + tag <- builtinOf (binaryId op) + _ <- builtinBinary @c @b @x tag -- A division is shareable only with a non-zero literal divisor. - if binaryName op `elem` ["div", "mod"] + if tag `elem` [BDiv, BMod] then case r of Lit (k :: b) | Just Refl <- testEquality (typeRep @b) (typeRep @Int) @@ -560,7 +562,7 @@ keyOf expr = fst <$> go maxKeyNodes expr else Just () (kl, afterLeft) <- go (budget - 1) l (kr, left) <- go afterLeft r - Just (KeyBinary (binaryName op) (SomeTypeRep (typeRep @c)) kl kr, left) + Just (KeyBinary tag (SomeTypeRep (typeRep @c)) kl kr, left) _ -> Nothing literalKey :: forall a. (Columnable a) => a -> Maybe Literal @@ -1181,7 +1183,7 @@ eval (GroupCtx gdf) expr@(Agg (FoldAgg _ Nothing (f :: a -> b -> a)) (Col name : Just col -> Flat <$> foldl1DirectGroups @b f col (valueIndices gdf) (offsets gdf) -- Fast path: countDistinct on a bare Col goes to the partitioned hash kernel. -eval ctx (Agg (MergeAgg "countDistinct" _ _ _ _) (Col name)) +eval ctx (Agg (MergeAgg (Builtin ACountDistinct) _ _ _ _) (Col name)) | Just Refl <- testEquality (typeRep @a) (typeRep @Int) , Just col <- getColumn name (ctxFrame ctx) , Just v <- distinctValue col = @@ -1238,7 +1240,7 @@ eval ctx expr@(Agg (FoldAgg _ (Just seed) (f :: a -> b -> a)) inner) = Group gs -> Flat . fromVector <$> V.mapM (foldlColumn @b @a f seed) gs -eval (FlatCtx df) (Agg (MergeAgg "mean" _ _ _ (_ :: acc -> a)) (Col name)) +eval (FlatCtx df) (Agg (MergeAgg (Builtin AMean) _ _ _ (_ :: acc -> a)) (Col name)) | Just Refl <- testEquality (typeRep @a) (typeRep @Double) , Just v <- fusedMean df name = Right (Scalar v) @@ -1265,7 +1267,7 @@ eval Flat . fromVector <$> V.mapM (fmap finalize . foldlColumn @b step seed) gs -- Fast path: whole-frame sum of a clean Int/Double column, one parallel loop. -eval (FlatCtx df) (Agg (FoldAgg "sum" Nothing (_ :: a -> b -> a)) inner) +eval (FlatCtx df) (Agg (FoldAgg (Builtin ASum) Nothing (_ :: a -> b -> a)) inner) | Just Refl <- testEquality (typeRep @a) (typeRep @b) , Just v <- fusedSum df inner = Right (Scalar v) @@ -1294,9 +1296,9 @@ eval ctx expr@(Agg (FoldAgg _ Nothing (f :: a -> b -> a)) inner) = fastUnaryFn :: forall b a. (Columnable b, Columnable a) => - T.Text -> (b -> a) -> (b -> a) + OpRef UnOpId -> (b -> a) -> (b -> a) fastUnaryFn name f - | name == "toDouble" + | Builtin UToDouble <- name , Just Refl <- testEquality (typeRep @a) (typeRep @Double) = integralToDouble @b f | otherwise = f @@ -1394,15 +1396,15 @@ fusedSum df expr = case expr of Binary op (Col name :: Expr c) (Lit (k :: b)) | Just Refl <- testEquality (typeRep @c) (typeRep @a) , Just Refl <- testEquality (typeRep @b) (typeRep @a) -> - withClean name (Just (binaryName op, k, False)) + builtinOf (binaryId op) >>= \tag -> withClean name (Just (tag, k, False)) Binary op (Lit (k :: c)) (Col name :: Expr b) | Just Refl <- testEquality (typeRep @c) (typeRep @a) , Just Refl <- testEquality (typeRep @b) (typeRep @a) -> - withClean name (Just (binaryName op, k, True)) + builtinOf (binaryId op) >>= \tag -> withClean name (Just (tag, k, True)) _ -> Nothing where - -- The operator is picked by name at a concrete element type. - withClean :: T.Text -> Maybe (T.Text, a, Bool) -> Maybe a + -- The operator is picked by tag at a concrete element type. + withClean :: T.Text -> Maybe (BinOpId, a, Bool) -> Maybe a withClean name combo = do UnboxedColumn Nothing (v :: VU.Vector t) <- getColumn name df Refl <- testEquality (typeRep @t) (typeRep @a) @@ -1412,20 +1414,20 @@ fusedSum df expr = case expr of Just Refl -> sumDouble combo v Nothing -> Nothing -sumInt :: Maybe (T.Text, Int, Bool) -> VU.Vector Int -> Maybe Int +sumInt :: Maybe (BinOpId, Int, Bool) -> VU.Vector Int -> Maybe Int sumInt combo v = case combo of Nothing -> Just (K.parSumBy id v) - Just ("add", k, _) -> Just (K.parSumBy (+ k) v) - Just ("mult", k, _) -> Just (K.parSumBy (* k) v) - Just ("sub", k, False) -> Just (K.parSumBy (subtract k) v) - Just ("sub", k, True) -> Just (K.parSumBy (k -) v) + Just (BAdd, k, _) -> Just (K.parSumBy (+ k) v) + Just (BMult, k, _) -> Just (K.parSumBy (* k) v) + Just (BSub, k, False) -> Just (K.parSumBy (subtract k) v) + Just (BSub, k, True) -> Just (K.parSumBy (k -) v) _ -> Nothing -sumDouble :: Maybe (T.Text, Double, Bool) -> VU.Vector Double -> Maybe Double +sumDouble :: Maybe (BinOpId, Double, Bool) -> VU.Vector Double -> Maybe Double sumDouble combo v = case combo of Nothing -> Just (K.parSumBy id v) - Just ("add", k, _) -> Just (K.parSumBy (+ k) v) - Just ("mult", k, _) -> Just (K.parSumBy (* k) v) - Just ("sub", k, False) -> Just (K.parSumBy (subtract k) v) - Just ("sub", k, True) -> Just (K.parSumBy (k -) v) + Just (BAdd, k, _) -> Just (K.parSumBy (+ k) v) + Just (BMult, k, _) -> Just (K.parSumBy (* k) v) + Just (BSub, k, False) -> Just (K.parSumBy (subtract k) v) + Just (BSub, k, True) -> Just (K.parSumBy (k -) v) _ -> Nothing diff --git a/dataframe-core/src-internal/DataFrame/Internal/Interpreter/Kernels.hs b/dataframe-core/src-internal/DataFrame/Internal/Interpreter/Kernels.hs index 7d91f57f..0f820f42 100644 --- a/dataframe-core/src-internal/DataFrame/Internal/Interpreter/Kernels.hs +++ b/dataframe-core/src-internal/DataFrame/Internal/Interpreter/Kernels.hs @@ -10,9 +10,9 @@ module DataFrame.Internal.Interpreter.Kernels ( Operand (..), Branch (..), CmpOp (..), - cmpOpByName, + cmpOpOf, ArithOp (..), - arithOpByName, + arithOpOf, compareOperands, arithOperands, zipOperands, @@ -44,6 +44,7 @@ import DataFrame.Internal.Data.PackedText ( selAt, sliceEqBytes, ) +import DataFrame.Internal.Expression.OpId (BinOpId (..)) import System.IO.Unsafe (unsafePerformIO) -- | One side of an operation: the same value in every row, or one per row. @@ -52,24 +53,24 @@ data Operand a = Literal !a | Column !(VU.Vector a) data CmpOp = CmpLt | CmpLeq | CmpGt | CmpGeq | CmpEq | CmpNeq deriving (Eq, Show) -cmpOpByName :: T.Text -> Maybe CmpOp -cmpOpByName name = case name of - "lt" -> Just CmpLt - "leq" -> Just CmpLeq - "gt" -> Just CmpGt - "geq" -> Just CmpGeq - "eq" -> Just CmpEq - "neq" -> Just CmpNeq +cmpOpOf :: BinOpId -> Maybe CmpOp +cmpOpOf op = case op of + BLt -> Just CmpLt + BLeq -> Just CmpLeq + BGt -> Just CmpGt + BGeq -> Just CmpGeq + BEq -> Just CmpEq + BNeq -> Just CmpNeq _ -> Nothing data ArithOp = ArithAdd | ArithSub | ArithMult deriving (Eq, Show) -arithOpByName :: T.Text -> Maybe ArithOp -arithOpByName name = case name of - "add" -> Just ArithAdd - "sub" -> Just ArithSub - "mult" -> Just ArithMult +arithOpOf :: BinOpId -> Maybe ArithOp +arithOpOf op = case op of + BAdd -> Just ArithAdd + BSub -> Just ArithSub + BMult -> Just ArithMult _ -> Nothing -- | Runs @k@ with the element type known, only for 'Double' or 'Int'. diff --git a/dataframe-expr-serializer/src/DataFrame/IR/ExprJson.hs b/dataframe-expr-serializer/src/DataFrame/IR/ExprJson.hs index 037f8191..a804165a 100644 --- a/dataframe-expr-serializer/src/DataFrame/IR/ExprJson.hs +++ b/dataframe-expr-serializer/src/DataFrame/IR/ExprJson.hs @@ -53,10 +53,17 @@ import qualified DataFrame.Functions as F import DataFrame.Internal.Column (Columnable) import DataFrame.Internal.Expression ( AggStrategy (..), - BinaryOp (binaryName), + BinOpId (..), + BinaryOp (binaryId), Expr (..), + OpRef (..), SomeExpr (..), - UnaryOp (unaryName), + UnOpId (..), + UnaryOp (unaryId), + aggName, + aggStrategyId, + binOpName, + unOpName, ) import qualified DataFrame.Internal.Expression as F import DataFrame.Internal.Expression.Reflect ( @@ -196,7 +203,7 @@ encodeExpr expr = case expr of Unary op (arg :: Expr b) -> do outTag <- requireTypeTag @a argTag <- requireTypeTag @b - opTag <- recognizeUnary (unaryName op) + opTag <- recognizeUnary (unaryId op) argEnc <- encodeExpr arg Right $ object @@ -209,7 +216,7 @@ encodeExpr expr = case expr of Binary op (lhs :: Expr c) (rhs :: Expr b) -> do outTag <- requireTypeTag @a argTag <- requireTypeTag @c - opTag <- recognizeBinary (binaryName op) + opTag <- recognizeBinary (binaryId op) lEnc <- encodeExpr lhs rEnc <- encodeExpr rhs Right $ @@ -244,11 +251,12 @@ encodeExpr expr = case expr of outTag <- requireTypeTag @a argTag <- requireTypeTag @b innerEnc <- encodeExpr inner + aggTag <- aggStrategyName strat Right $ object [ "node" .= ("agg" :: T.Text) , "out_type" .= outTag - , "agg" .= aggStrategyName strat + , "agg" .= aggTag , "arg_type" .= argTag , "arg" .= innerEnc ] @@ -271,10 +279,14 @@ encodeExpr expr = case expr of "DataFrame.IR.ExprJson.encodeExpr: unsupported type: " <> show (typeRep @x) -- | Wire-format name for an aggregation strategy. -aggStrategyName :: AggStrategy a b -> T.Text -aggStrategyName (CollectAgg n _) = n -aggStrategyName (FoldAgg n _ _) = n -aggStrategyName (MergeAgg n _ _ _ _) = n +aggStrategyName :: AggStrategy a b -> Either String T.Text +aggStrategyName strat = case aggStrategyId strat of + Builtin a -> Right (aggName a) + Named n -> Right n + Anonymous l -> + Left $ + "DataFrame.IR.ExprJson: anonymous aggregation (cannot serialize): " + <> T.unpack l -- | Encode an 'Expr' to a strict 'BS.ByteString' of JSON. encodeExprToBytes :: @@ -321,64 +333,81 @@ encodeLit v Left $ "DataFrame.IR.ExprJson.encodeLit: unsupported type: " <> show (typeRep @a) --- | Map a 'binaryName' string to a wire-format op name. -recognizeBinary :: T.Text -> Either String T.Text -recognizeBinary n = case n of - "add" -> Right "add" - "sub" -> Right "sub" - "mult" -> Right "mult" - "divide" -> Right "divide" - "eq" -> Right "eq" - "neq" -> Right "neq" - "lt" -> Right "lt" - "leq" -> Right "leq" - "gt" -> Right "gt" - "geq" -> Right "geq" - "and" -> Right "and" - "or" -> Right "or" - -- Nullable-aware aliases (lossless on non-null inputs) - "nulladd" -> Right "add" - "nullsub" -> Right "sub" - "nullmul" -> Right "mult" - "nulldiv" -> Right "divide" - "nulland" -> Right "and" - "nullor" -> Right "or" - "exponentiate" -> Right "exponentiate" - "logBase" -> Right "logBase" - "div" -> Right "div" - "mod" -> Right "mod" - "pow" -> Right "pow" - other -> +-- | Wire-format op name of a binary builtin. +recognizeBinary :: OpRef BinOpId -> Either String T.Text +recognizeBinary ident = case ident of + Builtin op -> case op of + BNullAdd -> Right "add" + BNullSub -> Right "sub" + BNullMul -> Right "mult" + BNullDiv -> Right "divide" + BNullAnd -> Right "and" + BNullOr -> Right "or" + BPowFrac -> Right "powFrac" + _ + | op `elem` wireBinary -> Right (binOpName op) + | otherwise -> unsupported (binOpName op) + Named n -> unsupported n + Anonymous l -> unsupported l + where + unsupported n = Left $ "DataFrame.IR.ExprJson: unsupported binary op (cannot serialize): " - <> T.unpack other + <> T.unpack n + wireBinary = + [ BAdd + , BSub + , BMult + , BDivide + , BEq + , BNeq + , BLt + , BLeq + , BGt + , BGeq + , BAnd + , BOr + , BExponentiate + , BLogBase + , BDiv + , BMod + , BPow + ] --- | Map a 'unaryName' string to a wire-format op name. -recognizeUnary :: T.Text -> Either String T.Text -recognizeUnary n = case n of - "not" -> Right "not" - "negate" -> Right "negate" - "abs" -> Right "abs" - "signum" -> Right "signum" - "exp" -> Right "exp" - "sqrt" -> Right "sqrt" - "log" -> Right "log" - "sin" -> Right "sin" - "cos" -> Right "cos" - "tan" -> Right "tan" - "asin" -> Right "asin" - "acos" -> Right "acos" - "atan" -> Right "atan" - "sinh" -> Right "sinh" - "cosh" -> Right "cosh" - "asinh" -> Right "asinh" - "acosh" -> Right "acosh" - "atanh" -> Right "atanh" - "toDouble" -> Right "toDouble" - other -> +-- | Wire-format op name of a unary builtin. +recognizeUnary :: OpRef UnOpId -> Either String T.Text +recognizeUnary ident = case ident of + Builtin op + | op `elem` wireUnary -> Right (unOpName op) + | otherwise -> unsupported (unOpName op) + Named n -> unsupported n + Anonymous l -> unsupported l + where + unsupported n = Left $ "DataFrame.IR.ExprJson: unsupported unary op (cannot serialize): " - <> T.unpack other + <> T.unpack n + wireUnary = + [ UNot + , UNegate + , UAbs + , USignum + , UExp + , USqrt + , ULog + , USin + , UCos + , UTan + , UAsin + , UAcos + , UAtan + , USinh + , UCosh + , UAsinh + , UAcosh + , UAtanh + , UToDouble + ] -- | Decode a JSON value into a 'SomeExpr' using the root @out_type@. decodeExprAny :: Aeson.Value -> Either String SomeExpr @@ -644,6 +673,11 @@ parseBinary op outType argType rawLhs rawRhs = case op of l <- parseExprAt @a rawLhs r <- parseExprAt @a rawRhs return $ SomeExpr (l ** r) + "logBase" -> withFloatingTypeTag outType $ \(_ :: Proxy a) -> do + requireSame outType argType + l <- parseExprAt @a rawLhs + r <- parseExprAt @a rawRhs + return $ SomeExpr (logBase l r) "nulladd" -> parseBinary "add" outType argType rawLhs rawRhs "nullsub" -> parseBinary "sub" outType argType rawLhs rawRhs "nullmul" -> parseBinary "mult" outType argType rawLhs rawRhs @@ -674,6 +708,12 @@ parseAgg aggName outType argType rawArg = case aggName of return $ SomeExpr (F.mean arg) "variance" -> realUnboxAgg F.variance "median" -> realUnboxAgg F.median + "stddev" -> realUnboxAgg F.stddev + "countDistinct" -> do + requireTag outType "int" + withOrdTypeTag argType $ \(_ :: Proxy a) -> do + arg <- parseExprAt @a rawArg + return $ SomeExpr (F.countDistinct arg) "collect" -> fail "DataFrame.IR.ExprJson: 'collect' is not supported in the wire format \ diff --git a/dataframe-lazy/src/DataFrame/Lazy/Internal/Executor.hs b/dataframe-lazy/src/DataFrame/Lazy/Internal/Executor.hs index 7d4437e2..ac105159 100644 --- a/dataframe-lazy/src/DataFrame/Lazy/Internal/Executor.hs +++ b/dataframe-lazy/src/DataFrame/Lazy/Internal/Executor.hs @@ -351,6 +351,13 @@ buildAggPlan aggs = foldl combine ([], [], id) (map processAgg aggs) 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 +377,10 @@ buildAggPlan aggs = foldl combine ([], [], id) (map processAgg aggs) [ ( 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 +394,10 @@ buildAggPlan aggs = foldl combine ([], [], id) (map processAgg aggs) [ ( 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 +424,13 @@ buildAggPlan aggs = foldl combine ([], [], id) (map processAgg aggs) 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 = diff --git a/dataframe-lazy/src/DataFrame/Lazy/Internal/Optimizer.hs b/dataframe-lazy/src/DataFrame/Lazy/Internal/Optimizer.hs index dcfdab85..f8b9363a 100644 --- a/dataframe-lazy/src/DataFrame/Lazy/Internal/Optimizer.hs +++ b/dataframe-lazy/src/DataFrame/Lazy/Internal/Optimizer.hs @@ -40,7 +40,7 @@ andExpr = E.Binary ( E.MkBinaryOp { E.binaryFn = (&&) - , E.binaryName = "and" + , E.binaryId = E.Builtin E.BAnd , E.binarySymbol = Just "&&" , E.binaryCommutative = True , E.binaryPrecedence = 3 diff --git a/dataframe-learn/src-internal/DataFrame/DecisionTree/CondVec.hs b/dataframe-learn/src-internal/DataFrame/DecisionTree/CondVec.hs index 7a4f9392..ff7bdf67 100644 --- a/dataframe-learn/src-internal/DataFrame/DecisionTree/CondVec.hs +++ b/dataframe-learn/src-internal/DataFrame/DecisionTree/CondVec.hs @@ -24,10 +24,12 @@ import qualified DataFrame.Functions as F import DataFrame.Internal.Column (TypedColumn (..), toVector) import DataFrame.Internal.DataFrame (DataFrame) import DataFrame.Internal.Expression ( - BinaryOp (binaryName), + BinOpId (..), + BinaryOp (binaryId), Expr (..), + builtinOf, eqExpr, - normalize, + structuralKey, ) import DataFrame.Internal.Interpreter (interpret) @@ -57,13 +59,14 @@ eitherToMaybe = either (const Nothing) Just -- | Truth vectors keyed by structural form, read-only once built. type CondCache = M.Map T.Text (VU.Vector Bool) --- | Structural key matching the candidate-dedup key. -condCacheKey :: Expr Bool -> T.Text -condCacheKey = T.pack . show . normalize +-- | Structural key; Nothing for conditions with anonymous operators. +condCacheKey :: Expr Bool -> Maybe T.Text +condCacheKey = fmap T.pack . structuralKey -- | Seed a cache from already-materialized candidate 'CondVec's (no interpret). condCacheFromVecs :: [CondVec] -> CondCache -condCacheFromVecs cvs = M.fromList [(condCacheKey (cvExpr cv), cvVec cv) | cv <- cvs] +condCacheFromVecs cvs = + M.fromList [(k, cvVec cv) | cv <- cvs, Just k <- [condCacheKey (cvExpr cv)]] -- | Add a tree's branch-condition vectors to a cache. addTreeCondsToCache :: DataFrame -> Tree a -> CondCache -> CondCache @@ -73,16 +76,16 @@ addTreeCondsToCache df = go go (Branch cond l r) c = go r (go l (insertCond df cond c)) insertCond :: DataFrame -> Expr Bool -> CondCache -> CondCache -insertCond df cond c - | M.member k c = c - | otherwise = - maybe c (\cv -> M.insert k (cvVec cv) c) (materializeCondVec df cond) - where - k = condCacheKey cond +insertCond df cond c = case condCacheKey cond of + Nothing -> c + Just k + | M.member k c -> c + | otherwise -> + maybe c (\cv -> M.insert k (cvVec cv) c) (materializeCondVec df cond) -- | A condition's truth vector: a cache hit, else interpret. lookupCondVec :: CondCache -> DataFrame -> Expr Bool -> Maybe (VU.Vector Bool) -lookupCondVec cache df cond = case M.lookup (condCacheKey cond) cache of +lookupCondVec cache df cond = case (`M.lookup` cache) =<< condCacheKey cond of hit@(Just _) -> hit Nothing -> cvVec <$> materializeCondVec df cond @@ -99,7 +102,7 @@ countErrorsByVec boolVals = length . filter misrouted -- | A same-column same-direction Double threshold comparison. data ThreshCmp = ThreshCmp { tcCol :: !T.Text - , tcName :: !T.Text + , tcOp :: !BinOpId , tcThr :: !Double , tcRebuild :: Double -> Expr Bool } @@ -108,29 +111,31 @@ asDoubleThreshold :: Expr Bool -> Maybe ThreshCmp asDoubleThreshold (Binary op (Col c :: Expr cc) (Lit (t :: tt))) = case ( testEquality (typeRep @cc) (typeRep @Double) , testEquality (typeRep @tt) (typeRep @Double) + , builtinOf (binaryId op) ) of - (Just Refl, Just Refl) -> Just (ThreshCmp c (binaryName op) t (Binary op (Col c) . Lit)) + (Just Refl, Just Refl, Just tag) -> + Just (ThreshCmp c tag t (Binary op (Col c) . Lit)) _ -> Nothing asDoubleThreshold _ = Nothing -directionalNames :: [T.Text] -directionalNames = ["lt", "leq", "gt", "geq"] +directionalOps :: [BinOpId] +directionalOps = [BLt, BLeq, BGt, BGeq] -- | Tighter (AND) or looser (OR) of two same-direction thresholds. -chooseThreshold :: Bool -> T.Text -> Double -> Double -> Double -chooseThreshold isAnd name t1 t2 +chooseThreshold :: Bool -> BinOpId -> Double -> Double -> Double +chooseThreshold isAnd op t1 t2 | leftDir = if isAnd then min t1 t2 else max t1 t2 | otherwise = if isAnd then max t1 t2 else min t1 t2 where - leftDir = name == "lt" || name == "leq" + leftDir = op == BLt || op == BLeq -- | Collapse two same-direction Double comparisons into one, if possible. consolidateThreshold :: Bool -> Expr Bool -> Expr Bool -> Maybe (Expr Bool) consolidateThreshold isAnd ea eb = do a <- asDoubleThreshold ea b <- asDoubleThreshold eb - if tcCol a == tcCol b && tcName a == tcName b && tcName a `elem` directionalNames - then Just (tcRebuild a (chooseThreshold isAnd (tcName a) (tcThr a) (tcThr b))) + if tcCol a == tcCol b && tcOp a == tcOp b && tcOp a `elem` directionalOps + then Just (tcRebuild a (chooseThreshold isAnd (tcOp a) (tcThr a) (tcThr b))) else Nothing -- | AND-combine two cached conditions, consolidating thresholds if possible. diff --git a/dataframe-learn/src-internal/DataFrame/DecisionTree/Pool.hs b/dataframe-learn/src-internal/DataFrame/DecisionTree/Pool.hs index 7d39b614..aff40591 100644 --- a/dataframe-learn/src-internal/DataFrame/DecisionTree/Pool.hs +++ b/dataframe-learn/src-internal/DataFrame/DecisionTree/Pool.hs @@ -36,7 +36,7 @@ import DataFrame.Internal.Expression ( eSize, eqExpr, getColumns, - normalize, + structuralKey, ) import Control.Parallel.Strategies (parListChunk, rdeepseq, using) @@ -157,12 +157,15 @@ admitKeys :: Set.Set String -> [CondVec] -> ([CondVec], Set.Set String) admitKeys = go [] where go acc seen [] = (reverse acc, seen) - go acc !seen (c : cs) - | structuralKey c `Set.member` seen = go acc seen cs - | otherwise = go (c : acc) (Set.insert (structuralKey c) seen) cs + go acc !seen (c : cs) = case cvKey c of + Nothing -> go (c : acc) seen cs + Just k + | k `Set.member` seen -> go acc seen cs + | otherwise -> go (c : acc) (Set.insert k seen) cs -structuralKey :: CondVec -> String -structuralKey = show . normalize . cvExpr +-- | Nothing for anonymous operators, which are never merged. +cvKey :: CondVec -> Maybe String +cvKey = structuralKey . cvExpr -- | Admit candidates by distinct truth vector, keeping the smallest expression. admitVecs :: @@ -187,17 +190,20 @@ dedupCVByExpr :: [CondVec] -> [CondVec] dedupCVByExpr = go Set.empty where go _ [] = [] - go seen (cv : cvs) - | structuralKey cv `Set.member` seen = go seen cvs - | otherwise = cv : go (Set.insert (structuralKey cv) seen) cvs + go seen (cv : cvs) = case cvKey cv of + Nothing -> cv : go seen cvs + Just k + | k `Set.member` seen -> go seen cvs + | otherwise -> cv : go (Set.insert k seen) cvs -- | Deduplicate expressions by normalized form, keeping the first. nubByExpr :: [Expr Bool] -> [Expr Bool] nubByExpr = go Set.empty where + go :: Set.Set String -> [Expr Bool] -> [Expr Bool] go _ [] = [] - go seen (e : es) - | k `Set.member` seen = go seen es - | otherwise = e : go (Set.insert k seen) es - where - k = show (normalize e) + go seen (e : es) = case structuralKey e of + Nothing -> e : go seen es + Just k + | k `Set.member` seen -> go seen es + | otherwise -> e : go (Set.insert k seen) es diff --git a/dataframe-operations/src-internal/DataFrame/Internal/Statistics.hs b/dataframe-operations/src-internal/DataFrame/Internal/Statistics.hs index ab99fa96..727744ff 100644 --- a/dataframe-operations/src-internal/DataFrame/Internal/Statistics.hs +++ b/dataframe-operations/src-internal/DataFrame/Internal/Statistics.hs @@ -14,6 +14,10 @@ import Control.Exception (throw) import Control.Monad.ST (runST) import DataFrame.Errors (DataFrameException (..)) +-- | The present values of a nullable vector, as doubles, in reverse order. +optionalToDoubleVector :: (Real a) => V.Vector (Maybe a) -> VU.Vector Double +optionalToDoubleVector = VU.fromList . V.foldl' (\acc e -> maybe acc ((: acc) . realToFrac) e) [] + mean' :: (Real a, VU.Unbox a) => VU.Vector a -> Double mean' samp | VU.null samp = throw $ EmptyDataSetException "mean" diff --git a/dataframe-operations/src/DataFrame/Functions.hs b/dataframe-operations/src/DataFrame/Functions.hs index 9f6f7be5..6ee75af4 100644 --- a/dataframe-operations/src/DataFrame/Functions.hs +++ b/dataframe-operations/src/DataFrame/Functions.hs @@ -27,7 +27,6 @@ import DataFrame.Internal.Expression import DataFrame.Internal.Statistics import Control.Applicative -import qualified Data.Char as Char import Data.Either import Data.Function (on) import Data.Int @@ -56,7 +55,8 @@ import Prelude as P lift :: (Columnable a, Columnable b) => (a -> b) -> Expr a -> Expr b lift f = - Unary (MkUnaryOp{unaryFn = f, unaryName = "unaryUdf", unarySymbol = Nothing}) + Unary + (MkUnaryOp{unaryFn = f, unaryId = Anonymous "unaryUdf", unarySymbol = Nothing}) lift2 :: (Columnable c, Columnable b, Columnable a) => @@ -68,7 +68,7 @@ lift2 f = Binary ( MkBinaryOp { binaryFn = f - , binaryName = "binaryUdf" + , binaryId = Anonymous "binaryUdf" , binarySymbol = Nothing , binaryCommutative = False , binaryPrecedence = 0 @@ -83,7 +83,12 @@ nullLift :: Expr (NullLift1Result a r) nullLift f = Unary - (MkUnaryOp{unaryFn = applyNull1 f, unaryName = "nullLift", unarySymbol = Nothing}) + ( MkUnaryOp + { unaryFn = applyNull1 f + , unaryId = Anonymous "nullLift" + , unarySymbol = Nothing + } + ) -- | Lift a binary function over nullable or non-nullable column expressions. nullLift2 :: @@ -96,7 +101,7 @@ nullLift2 f = Binary ( MkBinaryOp { binaryFn = applyNull2 f - , binaryName = "nullLift2" + , binaryId = Anonymous "nullLift2" , binarySymbol = Nothing , binaryCommutative = False , binaryPrecedence = 0 @@ -105,24 +110,24 @@ nullLift2 f = -- | Lenient coercion of a column to @Maybe a@; failures become 'Nothing'. cast :: forall a. (Columnable a, Read a) => T.Text -> Expr (Maybe a) -cast colName = CastWith colName "cast" (either (const Nothing) Just) +cast colName = CastWith colName (Builtin CCast) (either (const Nothing) Just) -- | Lenient coercion that substitutes a default for unconvertible values. castWithDefault :: forall a. (Columnable a, Read a) => a -> T.Text -> Expr a castWithDefault def colName = - CastWith colName ("castWithDefault:" <> T.pack (show def)) (fromRight def) + CastWith colName (Builtin (CCastWithDefault (T.pack (show def)))) (fromRight def) -- | Lenient coercion to @Either T.Text a@, keeping unparsed values as 'Left'. castEither :: forall a. (Columnable a, Read a) => T.Text -> Expr (Either T.Text a) -castEither colName = CastWith colName "castEither" (either (Left . T.pack) Right) +castEither colName = CastWith colName (Builtin CCastEither) (either (Left . T.pack) Right) -- | Lenient coercion that crashes at evaluation if a @Nothing@ is found. unsafeCast :: forall a. (Columnable a, Read a) => T.Text -> Expr a unsafeCast colName = CastWith colName - "unsafeCast" + (Builtin CUnsafeCast) (fromRight (error "unsafeCast: unexpected Nothing in column")) castExpr :: @@ -130,13 +135,13 @@ castExpr :: (Columnable b, Columnable src, Read b) => Expr src -> Expr (Maybe b) -castExpr = CastExprWith @b @(Maybe b) @src "castExpr" (either (const Nothing) Just) +castExpr = CastExprWith @b @(Maybe b) @src (Builtin CCast) (either (const Nothing) Just) castExprWithDefault :: forall b src. (Columnable b, Columnable src, Read b) => b -> Expr src -> Expr b castExprWithDefault def = CastExprWith @b @b @src - ("castExprWithDefault:" <> T.pack (show def)) + (Builtin (CCastWithDefault (T.pack (show def)))) (fromRight def) castExprEither :: @@ -146,14 +151,14 @@ castExprEither :: Expr (Either T.Text b) castExprEither = CastExprWith @b @(Either T.Text b) @src - "castExprEither" + (Builtin CCastEither) (either (Left . T.pack) Right) unsafeCastExpr :: forall b src. (Columnable b, Columnable src, Read b) => Expr src -> Expr b unsafeCastExpr = CastExprWith @b @b @src - "unsafeCastExpr" + (Builtin CUnsafeCast) (fromRight (error "unsafeCastExpr: unexpected Nothing in column")) toDouble :: (Columnable a, Real a) => Expr a -> Expr Double @@ -161,34 +166,34 @@ toDouble = Unary ( MkUnaryOp { unaryFn = realToFrac - , unaryName = "toDouble" + , unaryId = Builtin UToDouble , unarySymbol = Nothing } ) infix 8 `div` div :: (Integral a, Columnable a) => Expr a -> Expr a -> Expr a -div = lift2Decorated Prelude.div "div" (Just "//") False 7 +div = lift2Builtin BDiv Prelude.div (Just "//") False 7 mod :: (Integral a, Columnable a) => Expr a -> Expr a -> Expr a -mod = lift2Decorated Prelude.mod "mod" Nothing False 7 +mod = lift2Builtin BMod Prelude.mod Nothing False 7 eq :: (Columnable a, Eq a, a ~ BaseType a) => Expr a -> Expr a -> Expr Bool -eq = lift2Decorated (==) "eq" (Just "==") True 4 +eq = lift2Builtin BEq (==) (Just "==") True 4 lt :: (Columnable a, Ord a, a ~ BaseType a) => Expr a -> Expr a -> Expr Bool -lt = lift2Decorated (<) "lt" (Just "<") False 4 +lt = lift2Builtin BLt (<) (Just "<") False 4 gt :: (Columnable a, Ord a, a ~ BaseType a) => Expr a -> Expr a -> Expr Bool -gt = lift2Decorated (>) "gt" (Just ">") False 4 +gt = lift2Builtin BGt (>) (Just ">") False 4 leq :: (Columnable a, Ord a, Eq a, a ~ BaseType a) => Expr a -> Expr a -> Expr Bool -leq = lift2Decorated (<=) "leq" (Just "<=") False 4 +leq = lift2Builtin BLeq (<=) (Just "<=") False 4 geq :: (Columnable a, Ord a, Eq a, a ~ BaseType a) => Expr a -> Expr a -> Expr Bool -geq = lift2Decorated (>=) "geq" (Just ">=") False 4 +geq = lift2Builtin BGeq (>=) (Just ">=") False 4 and :: Expr Bool -> Expr Bool -> Expr Bool and = (.&&) @@ -199,10 +204,10 @@ or = (.||) not :: Expr Bool -> Expr Bool not = Unary - (MkUnaryOp{unaryFn = Prelude.not, unaryName = "not", unarySymbol = Just "~"}) + (MkUnaryOp{unaryFn = Prelude.not, unaryId = Builtin UNot, unarySymbol = Just "~"}) count :: (Columnable a) => Expr a -> Expr Int -count = Agg (MergeAgg "count" (0 :: Int) (\c _ -> c + 1) (+) id) +count = Agg (MergeAgg (Builtin ACount) (0 :: Int) (\c _ -> c + 1) (+) id) {-# SPECIALIZE count :: Expr Double -> Expr Int #-} {-# SPECIALIZE count :: Expr Float -> Expr Int #-} {-# SPECIALIZE count :: Expr Int -> Expr Int #-} @@ -217,7 +222,7 @@ countDistinct :: forall a. (Columnable a, Ord a) => Expr a -> Expr Int countDistinct = Agg ( MergeAgg - "countDistinct" + (Builtin ACountDistinct) (S.empty :: S.Set a) (flip S.insert) S.union @@ -237,7 +242,7 @@ countAll = count (Lit (0 :: Int)) {-# INLINE countAll #-} collect :: (Columnable a) => Expr a -> Expr [a] -collect = Agg (FoldAgg "collect" (Just []) (flip (:))) +collect = Agg (FoldAgg (Builtin ACollect) (Just []) (flip (:))) {-# SPECIALIZE collect :: Expr Double -> Expr [Double] #-} {-# SPECIALIZE collect :: Expr Float -> Expr [Float] #-} {-# SPECIALIZE collect :: Expr Int -> Expr [Int] #-} @@ -247,7 +252,7 @@ mode :: (Ord a, Columnable a, Eq a) => Expr a -> Expr a mode = Agg ( CollectAgg - "mode" + (Builtin AMode) ( fst . L.maximumBy (compare `on` snd) . M.toList @@ -264,7 +269,7 @@ mode = {-# INLINEABLE mode #-} minimum :: (Columnable a, Ord a) => Expr a -> Expr a -minimum = Agg (FoldAgg "minimum" Nothing Prelude.min) +minimum = Agg (FoldAgg (Builtin AMinimum) Nothing Prelude.min) {-# SPECIALIZE DataFrame.Functions.minimum :: Expr Double -> Expr Double #-} {-# SPECIALIZE DataFrame.Functions.minimum :: Expr Float -> Expr Float #-} {-# SPECIALIZE DataFrame.Functions.minimum :: Expr Int -> Expr Int #-} @@ -275,7 +280,7 @@ minimum = Agg (FoldAgg "minimum" Nothing Prelude.min) {-# INLINEABLE DataFrame.Functions.minimum #-} maximum :: (Columnable a, Ord a) => Expr a -> Expr a -maximum = Agg (FoldAgg "maximum" Nothing Prelude.max) +maximum = Agg (FoldAgg (Builtin AMaximum) Nothing Prelude.max) {-# SPECIALIZE DataFrame.Functions.maximum :: Expr Double -> Expr Double #-} {-# SPECIALIZE DataFrame.Functions.maximum :: Expr Float -> Expr Float #-} {-# SPECIALIZE DataFrame.Functions.maximum :: Expr Int -> Expr Int #-} @@ -286,7 +291,7 @@ maximum = Agg (FoldAgg "maximum" Nothing Prelude.max) {-# INLINEABLE DataFrame.Functions.maximum #-} sum :: forall a. (Columnable a, Num a) => Expr a -> Expr a -sum = Agg (FoldAgg "sum" Nothing (+)) +sum = Agg (FoldAgg (Builtin ASum) Nothing (+)) {-# SPECIALIZE DataFrame.Functions.sum :: Expr Double -> Expr Double #-} {-# SPECIALIZE DataFrame.Functions.sum :: Expr Float -> Expr Float #-} {-# SPECIALIZE DataFrame.Functions.sum :: Expr Int -> Expr Int #-} @@ -297,7 +302,7 @@ sum = Agg (FoldAgg "sum" Nothing (+)) {-# INLINEABLE DataFrame.Functions.sum #-} sumMaybe :: forall a. (Columnable a, Num a) => Expr (Maybe a) -> Expr a -sumMaybe = Agg (CollectAgg "sumMaybe" (P.sum . Maybe.catMaybes . V.toList)) +sumMaybe = Agg (CollectAgg (Builtin ASumMaybe) (P.sum . Maybe.catMaybes . V.toList)) {-# SPECIALIZE sumMaybe :: Expr (Maybe Double) -> Expr Double #-} {-# SPECIALIZE sumMaybe :: Expr (Maybe Float) -> Expr Float #-} {-# SPECIALIZE sumMaybe :: Expr (Maybe Int) -> Expr Int #-} @@ -311,7 +316,7 @@ mean :: (Columnable a, Real a) => Expr a -> Expr Double mean = Agg ( MergeAgg - "mean" + (Builtin AMean) ((0.0, 0) :: (Double, Int)) (\(!s, !c) x -> (s + realToFrac x, c + 1)) (\(!s1, !c1) (!s2, !c2) -> (s1 + s2, c1 + c2)) @@ -328,29 +333,27 @@ mean = -- | The k largest values per group. topK :: (Columnable a, Ord a) => Int -> Expr a -> Expr [a] -topK = kExtremes "topK" (>) +topK = kExtremes True {-# INLINEABLE topK #-} -- | The k smallest values per group. bottomK :: (Columnable a, Ord a) => Int -> Expr a -> Expr [a] -bottomK = kExtremes "bottomK" (<) +bottomK = kExtremes False {-# INLINEABLE bottomK #-} --- | Top k values after applying an ordering function, dropping NaNs. -kExtremes :: - forall a. - (Columnable a, Ord a) => - T.Text -> (a -> a -> Bool) -> Int -> Expr a -> Expr [a] -kExtremes opName wins k = +-- | The k largest (True) or smallest (False) values, dropping NaNs. +kExtremes :: forall a. (Columnable a, Ord a) => Bool -> Int -> Expr a -> Expr [a] +kExtremes largest k = Agg ( MergeAgg - (opName <> "_" <> T.pack (show k)) + (Builtin (if largest then ATopK k else ABottomK k)) [] step (L.foldl' step) id ) where + wins = if largest then (>) else (<) step acc x | isNaNLike x = acc | otherwise = insert x acc @@ -366,7 +369,7 @@ kExtremes opName wins k = | otherwise = False meanMaybe :: forall a. (Columnable a, Real a) => Expr (Maybe a) -> Expr Double -meanMaybe = Agg (CollectAgg "meanMaybe" (mean' . optionalToDoubleVector)) +meanMaybe = Agg (CollectAgg (Builtin AMeanMaybe) (mean' . optionalToDoubleVector)) {-# SPECIALIZE meanMaybe :: Expr (Maybe Double) -> Expr Double #-} {-# SPECIALIZE meanMaybe :: Expr (Maybe Float) -> Expr Double #-} {-# SPECIALIZE meanMaybe :: Expr (Maybe Int) -> Expr Double #-} @@ -377,7 +380,7 @@ meanMaybe = Agg (CollectAgg "meanMaybe" (mean' . optionalToDoubleVector)) {-# INLINEABLE meanMaybe #-} variance :: (Columnable a, Real a, VU.Unbox a) => Expr a -> Expr Double -variance = Agg (CollectAgg "variance" variance') +variance = Agg (CollectAgg (Builtin AVariance) variance') {-# SPECIALIZE variance :: Expr Double -> Expr Double #-} {-# SPECIALIZE variance :: Expr Float -> Expr Double #-} {-# SPECIALIZE variance :: Expr Int -> Expr Double #-} @@ -388,7 +391,7 @@ variance = Agg (CollectAgg "variance" variance') {-# INLINEABLE variance #-} median :: (Columnable a, Real a, VU.Unbox a) => Expr a -> Expr Double -median = Agg (CollectAgg "median" median') +median = Agg (CollectAgg (Builtin AMedian) median') {-# SPECIALIZE median :: Expr Double -> Expr Double #-} {-# SPECIALIZE median :: Expr Float -> Expr Double #-} {-# SPECIALIZE median :: Expr Int -> Expr Double #-} @@ -399,7 +402,7 @@ median = Agg (CollectAgg "median" median') {-# INLINEABLE median #-} medianMaybe :: (Columnable a, Real a) => Expr (Maybe a) -> Expr Double -medianMaybe = Agg (CollectAgg "meanMaybe" (median' . optionalToDoubleVector)) +medianMaybe = Agg (CollectAgg (Builtin AMedianMaybe) (median' . optionalToDoubleVector)) {-# SPECIALIZE medianMaybe :: Expr (Maybe Double) -> Expr Double #-} {-# SPECIALIZE medianMaybe :: Expr (Maybe Float) -> Expr Double #-} {-# SPECIALIZE medianMaybe :: Expr (Maybe Int) -> Expr Double #-} @@ -409,23 +412,16 @@ medianMaybe = Agg (CollectAgg "meanMaybe" (median' . optionalToDoubleVector)) {-# SPECIALIZE medianMaybe :: Expr (Maybe Int64) -> Expr Double #-} {-# INLINEABLE medianMaybe #-} -optionalToDoubleVector :: (Real a) => V.Vector (Maybe a) -> VU.Vector Double -optionalToDoubleVector = - VU.fromList - . V.foldl' - (\acc e -> if Maybe.isJust e then realToFrac (Maybe.fromMaybe 0 e) : acc else acc) - [] - percentile :: Int -> Expr Double -> Expr Double percentile n = Agg ( CollectAgg - (T.pack $ "percentile " ++ show n) + (Builtin (APercentile n)) (percentile' n) ) stddev :: (Columnable a, Real a, VU.Unbox a) => Expr a -> Expr Double -stddev = Agg (CollectAgg "stddev" (sqrt . variance')) +stddev = Agg (CollectAgg (Builtin AStddev) (sqrt . variance')) {-# SPECIALIZE stddev :: Expr Double -> Expr Double #-} {-# SPECIALIZE stddev :: Expr Float -> Expr Double #-} {-# SPECIALIZE stddev :: Expr Int -> Expr Double #-} @@ -436,7 +432,9 @@ stddev = Agg (CollectAgg "stddev" (sqrt . variance')) {-# INLINEABLE stddev #-} stddevMaybe :: forall a. (Columnable a, Real a) => Expr (Maybe a) -> Expr Double -stddevMaybe = Agg (CollectAgg "stddevMaybe" (sqrt . variance' . optionalToDoubleVector)) +stddevMaybe = + Agg + (CollectAgg (Builtin AStddevMaybe) (sqrt . variance' . optionalToDoubleVector)) {-# SPECIALIZE stddevMaybe :: Expr (Maybe Double) -> Expr Double #-} {-# SPECIALIZE stddevMaybe :: Expr (Maybe Float) -> Expr Double #-} {-# SPECIALIZE stddevMaybe :: Expr (Maybe Int) -> Expr Double #-} @@ -450,21 +448,21 @@ zScore :: Expr Double -> Expr Double zScore c = (c - mean c) / stddev c pow :: (Columnable a, Num a) => Expr a -> Int -> Expr a -pow expr i = lift2Decorated (^) "pow" (Just "^") False 8 expr (Lit i) +pow expr i = lift2Builtin BPow (^) (Just "^") False 8 expr (Lit i) {-# SPECIALIZE pow :: Expr Double -> Int -> Expr Double #-} {-# SPECIALIZE pow :: Expr Float -> Int -> Expr Float #-} {-# SPECIALIZE pow :: Expr Int -> Int -> Expr Int #-} {-# INLINEABLE pow #-} relu :: (Columnable a, Num a, Ord a) => Expr a -> Expr a -relu = liftDecorated (Prelude.max 0) "relu" Nothing +relu = liftBuiltin URelu (Prelude.max 0) Nothing {-# SPECIALIZE relu :: Expr Double -> Expr Double #-} {-# SPECIALIZE relu :: Expr Float -> Expr Float #-} {-# SPECIALIZE relu :: Expr Int -> Expr Int #-} {-# INLINEABLE relu #-} min :: (Columnable a, Ord a) => Expr a -> Expr a -> Expr a -min = lift2Decorated Prelude.min "min" Nothing True 1 +min = lift2Builtin BMin Prelude.min Nothing True 1 {-# SPECIALIZE DataFrame.Functions.min :: Expr Double -> Expr Double -> Expr Double #-} @@ -475,7 +473,7 @@ min = lift2Decorated Prelude.min "min" Nothing True 1 {-# INLINEABLE DataFrame.Functions.min #-} max :: (Columnable a, Ord a) => Expr a -> Expr a -> Expr a -max = lift2Decorated Prelude.max "max" Nothing True 1 +max = lift2Builtin BMax Prelude.max Nothing True 1 {-# SPECIALIZE DataFrame.Functions.max :: Expr Double -> Expr Double -> Expr Double #-} @@ -492,31 +490,31 @@ reduce :: a -> (a -> b -> a) -> Expr a -reduce expr start f = Agg (FoldAgg "foldUdf" (Just start) f) expr +reduce expr start f = Agg (FoldAgg (Anonymous "foldUdf") (Just start) f) expr {-# INLINEABLE reduce #-} toMaybe :: (Columnable a) => Expr a -> Expr (Maybe a) -toMaybe = liftDecorated Just "toMaybe" Nothing +toMaybe = liftBuiltin UToMaybe Just Nothing {-# SPECIALIZE toMaybe :: Expr Double -> Expr (Maybe Double) #-} {-# SPECIALIZE toMaybe :: Expr Float -> Expr (Maybe Float) #-} {-# SPECIALIZE toMaybe :: Expr Int -> Expr (Maybe Int) #-} {-# INLINEABLE toMaybe #-} fromMaybe :: (Columnable a) => a -> Expr (Maybe a) -> Expr a -fromMaybe d = liftDecorated (Maybe.fromMaybe d) "fromMaybe" Nothing +fromMaybe d = liftBuiltin (UFromMaybe (T.pack (show d))) (Maybe.fromMaybe d) Nothing {-# SPECIALIZE fromMaybe :: Double -> Expr (Maybe Double) -> Expr Double #-} {-# SPECIALIZE fromMaybe :: Float -> Expr (Maybe Float) -> Expr Float #-} {-# SPECIALIZE fromMaybe :: Int -> Expr (Maybe Int) -> Expr Int #-} {-# INLINEABLE fromMaybe #-} isJust :: (Columnable a) => Expr (Maybe a) -> Expr Bool -isJust = liftDecorated Maybe.isJust "isJust" Nothing +isJust = liftBuiltin UIsJust Maybe.isJust Nothing {-# SPECIALIZE isJust :: Expr (Maybe Double) -> Expr Bool #-} {-# SPECIALIZE isJust :: Expr (Maybe Int) -> Expr Bool #-} {-# INLINEABLE isJust #-} isNothing :: (Columnable a) => Expr (Maybe a) -> Expr Bool -isNothing = liftDecorated Maybe.isNothing "isNothing" Nothing +isNothing = liftBuiltin UIsNothing Maybe.isNothing Nothing {-# SPECIALIZE isNothing :: Expr (Maybe Double) -> Expr Bool #-} {-# SPECIALIZE isNothing :: Expr (Maybe Int) -> Expr Bool #-} {-# INLINEABLE isNothing #-} @@ -532,7 +530,7 @@ isNotNull = isJust {-# INLINEABLE isNotNull #-} fromJust :: (Columnable a) => Expr (Maybe a) -> Expr a -fromJust = liftDecorated Maybe.fromJust "fromJust" Nothing +fromJust = liftBuiltin UFromJust Maybe.fromJust Nothing {-# SPECIALIZE fromJust :: Expr (Maybe Double) -> Expr Double #-} {-# SPECIALIZE fromJust :: Expr (Maybe Int) -> Expr Int #-} {-# INLINEABLE fromJust #-} @@ -543,7 +541,7 @@ whenPresent :: (a -> b) -> Expr (Maybe a) -> Expr (Maybe b) -whenPresent f = liftDecorated (fmap f) "whenPresent" Nothing +whenPresent f = Unary (MkUnaryOp (fmap f) (Anonymous "whenPresent") Nothing) {-# INLINEABLE whenPresent #-} whenBothPresent :: @@ -553,7 +551,9 @@ whenBothPresent :: Expr (Maybe a) -> Expr (Maybe b) -> Expr (Maybe c) -whenBothPresent f = lift2Decorated (\l r -> f <$> l <*> r) "whenBothPresent" Nothing False 0 +whenBothPresent f = + Binary + (MkBinaryOp (\l r -> f <$> l <*> r) (Anonymous "whenBothPresent") Nothing False 0) {-# INLINEABLE whenBothPresent #-} recode :: @@ -566,7 +566,7 @@ recode mapping = Unary ( MkUnaryOp { unaryFn = (`lookup` mapping) - , unaryName = "recode " <> T.pack (show mapping) + , unaryId = Builtin (URecode (T.pack (show mapping))) , unarySymbol = Nothing } ) @@ -592,57 +592,57 @@ recodeWithDefault d mapping = Unary ( MkUnaryOp { unaryFn = Maybe.fromMaybe d . (`lookup` mapping) - , unaryName = - "recodeWithDefault " <> T.pack (show d) <> " " <> T.pack (show mapping) + , unaryId = + Builtin (URecodeWithDefault (T.pack (show d)) (T.pack (show mapping))) , unarySymbol = Nothing } ) firstOrNothing :: (Columnable a) => Expr [a] -> Expr (Maybe a) -firstOrNothing = liftDecorated Maybe.listToMaybe "firstOrNothing" Nothing +firstOrNothing = liftBuiltin UFirstOrNothing Maybe.listToMaybe Nothing lastOrNothing :: (Columnable a) => Expr [a] -> Expr (Maybe a) -lastOrNothing = liftDecorated (Maybe.listToMaybe . reverse) "lastOrNothing" Nothing +lastOrNothing = liftBuiltin ULastOrNothing (Maybe.listToMaybe . reverse) Nothing -- | Whether the text contains the literal substring. contains :: T.Text -> Expr T.Text -> Expr Bool contains needle expr = - lift2Decorated (flip T.isInfixOf) "contains" Nothing False 4 expr (Lit needle) + lift2Builtin BContains (flip T.isInfixOf) Nothing False 4 expr (Lit needle) -- | Character count. length :: Expr T.Text -> Expr Int -length = liftDecorated T.length "textLength" Nothing +length = liftBuiltin UTextLength T.length Nothing splitOn :: T.Text -> Expr T.Text -> Expr [T.Text] -splitOn delim = liftDecorated (T.splitOn delim) "splitOn" Nothing +splitOn delim = liftBuiltin (USplitOn delim) (T.splitOn delim) Nothing match :: T.Text -> Expr T.Text -> Expr (Maybe T.Text) match regex = - liftDecorated + liftBuiltin + (UMatch regex) ((\r -> if T.null r then Nothing else Just r) . (=~ regex)) - ("match " <> T.pack (show regex)) Nothing matchAll :: T.Text -> Expr T.Text -> Expr [T.Text] matchAll regex = - liftDecorated + liftBuiltin + (UMatchAll regex) (getAllTextMatches . (=~ regex)) - ("matchAll " <> T.pack (show regex)) Nothing parseDate :: (ParseTime t, Columnable t) => T.Text -> Expr T.Text -> Expr (Maybe t) parseDate format = - liftDecorated + liftBuiltin + (UParseDate format) (parseTimeM True defaultTimeLocale (T.unpack format) . T.unpack) - ("parseDate " <> format) Nothing daysBetween :: Expr Day -> Expr Day -> Expr Int daysBetween = - lift2Decorated + lift2Builtin + BDaysBetween (\d1 d2 -> fromIntegral (diffDays d1 d2)) - "daysBetween" Nothing False 2 @@ -653,71 +653,8 @@ bind :: (a -> m b) -> Expr (m a) -> Expr (m b) -bind f = liftDecorated (>>= f) "bind" Nothing +bind f = Unary (MkUnaryOp (>>= f) (Anonymous "bind") Nothing) -- | Window function: evaluate an expression partitioned by the given columns. over :: (Columnable a) => [T.Text] -> Expr a -> Expr a over = Over - --- Haskell Report 2.4: https://www.haskell.org/definition/haskell2010.pdf -isReservedId :: T.Text -> Bool -isReservedId t = case t of - "case" -> True - "class" -> True - "data" -> True - "default" -> True - "deriving" -> True - "do" -> True - "else" -> True - "foreign" -> True - "if" -> True - "import" -> True - "in" -> True - "infix" -> True - "infixl" -> True - "infixr" -> True - "instance" -> True - "let" -> True - "module" -> True - "newtype" -> True - "of" -> True - "then" -> True - "type" -> True - "where" -> True - _ -> False - -isVarId :: T.Text -> Bool -isVarId t = case T.uncons t of - -- '_' is skipped on purpose so "Name" and "_Name_" cannot collide. - Just (c, _) -> Char.isLower c && Char.isAlpha c - Nothing -> False - -isHaskellIdentifier :: T.Text -> Bool -isHaskellIdentifier t = Prelude.not (isVarId t) || isReservedId t - -sanitize :: T.Text -> T.Text -sanitize t - | isValid = t - | isHaskellIdentifier t' = "_" <> t' <> "_" - | otherwise = t' - where - isValid = - Prelude.not (isHaskellIdentifier t) - && isVarId t - && T.all Char.isAlphaNum t - t' = T.map replaceInvalidCharacters . T.filter (Prelude.not . parentheses) $ t - replaceInvalidCharacters c - | Char.isUpper c = Char.toLower c - | Char.isSpace c = '_' - | Char.isPunctuation c = '_' -- '-' will also become a '_' - | Char.isSymbol c = '_' - | Char.isAlphaNum c = c -- Blanket condition - | otherwise = '_' -- If we're unsure we'll default to an underscore - parentheses c = case c of - '(' -> True - ')' -> True - '{' -> True - '}' -> True - '[' -> True - ']' -> True - _ -> False diff --git a/dataframe-operations/src/DataFrame/Operations/Statistics.hs b/dataframe-operations/src/DataFrame/Operations/Statistics.hs index 6f32dce8..d50b957f 100644 --- a/dataframe-operations/src/DataFrame/Operations/Statistics.hs +++ b/dataframe-operations/src/DataFrame/Operations/Statistics.hs @@ -235,13 +235,6 @@ _getColumnAsDouble name df = case getColumn name df of _ -> Nothing {-# INLINE _getColumnAsDouble #-} -optionalToDoubleVector :: (Real a) => V.Vector (Maybe a) -> VU.Vector Double -optionalToDoubleVector = - VU.fromList - . V.foldl' - (\acc e -> if isJust e then realToFrac (fromMaybe 0 e) : acc else acc) - [] - -- | Calculates the sum of a given column as a standalone value. sum :: forall a. (Columnable a, Num a) => Expr a -> DataFrame -> a diff --git a/dataframe-operations/src/DataFrame/Typed/Expr.hs b/dataframe-operations/src/DataFrame/Typed/Expr.hs index 433ee4fc..fbedf98e 100644 --- a/dataframe-operations/src/DataFrame/Typed/Expr.hs +++ b/dataframe-operations/src/DataFrame/Typed/Expr.hs @@ -105,7 +105,6 @@ module DataFrame.Typed.Expr ( module DataFrame.Typed.Expr.Extra, ) where -import Data.Either (fromRight) import Data.Proxy (Proxy (..)) import Data.String (IsString (..)) import qualified Data.Text as T @@ -116,15 +115,16 @@ import qualified DataFrame.Functions as F import DataFrame.Internal.Column (Columnable) import DataFrame.Internal.Column.Types (Promote, PromoteDiv) import DataFrame.Internal.Expression ( + BinOpId (..), BinUDF (..), Expr (..), - UnUDF (..), + OpRef (..), ) import DataFrame.Internal.Expression.Operators.Nullable ( BaseType, DivWidenOp, NullCmpResult, - NullLift1Op (applyNull1), + NullLift1Op, NullLift1Result, NullLift2Op (applyNull2), NullLift2Result, @@ -229,13 +229,13 @@ instance (IsString a, Columnable a) => IsString (TExpr cols a) where -- | Lift a unary function into a typed expression. lift :: (Columnable a, Columnable b) => (a -> b) -> TExpr cols a -> TExpr cols b -lift f (TExpr e) = TExpr (Unary (MkUnaryOp f "unaryUdf" Nothing) e) +lift f (TExpr e) = TExpr (F.lift f e) -- | Lift a binary function into typed expressions. lift2 :: (Columnable a, Columnable b, Columnable c) => (a -> b -> c) -> TExpr cols a -> TExpr cols b -> TExpr cols c -lift2 f (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp f "binaryUdf" Nothing False 0) a b) +lift2 f (TExpr a) (TExpr b) = TExpr (F.lift2 f a b) -- | Typed 'nullLift': lift a unary function with nullable propagation. nullLift :: @@ -243,7 +243,7 @@ nullLift :: (BaseType a -> r) -> TExpr cols a -> TExpr cols (NullLift1Result a r) -nullLift f (TExpr e) = TExpr (Unary (MkUnaryOp (applyNull1 f) "nullLift" Nothing) e) +nullLift f (TExpr e) = TExpr (F.nullLift f e) -- | Typed 'nullLift2': lift a binary function with nullable propagation. nullLift2 :: @@ -252,7 +252,7 @@ nullLift2 :: TExpr cols a -> TExpr cols b -> TExpr cols (NullLift2Result a b r) -nullLift2 f (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (applyNull2 f) "nullLift2" Nothing False 0) a b) +nullLift2 f (TExpr a) (TExpr b) = TExpr (F.nullLift2 f a b) infixl 4 .==., ./=., .<., .<=., .>=., .>. infix 4 .==, ./=, .<, .<=, .>=, .> @@ -264,27 +264,27 @@ infixr 8 .^^., .^^, .^., .^ (.==.) :: (Columnable a, Eq a) => TExpr cols a -> TExpr cols a -> TExpr cols Bool -(.==.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (==) "eq" (Just "==") True 4) a b) +(.==.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (==) (Builtin BEq) (Just "==") True 4) a b) (./=.) :: (Columnable a, Eq a) => TExpr cols a -> TExpr cols a -> TExpr cols Bool -(./=.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (/=) "neq" (Just "/=") True 4) a b) +(./=.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (/=) (Builtin BNeq) (Just "/=") True 4) a b) (.<.) :: (Columnable a, Ord a) => TExpr cols a -> TExpr cols a -> TExpr cols Bool -(.<.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (<) "lt" (Just "<") False 4) a b) +(.<.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (<) (Builtin BLt) (Just "<") False 4) a b) (.<=.) :: (Columnable a, Ord a) => TExpr cols a -> TExpr cols a -> TExpr cols Bool -(.<=.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (<=) "leq" (Just "<=") False 4) a b) +(.<=.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (<=) (Builtin BLeq) (Just "<=") False 4) a b) (.>=.) :: (Columnable a, Ord a) => TExpr cols a -> TExpr cols a -> TExpr cols Bool -(.>=.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (>=) "geq" (Just ">=") False 4) a b) +(.>=.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (>=) (Builtin BGeq) (Just ">=") False 4) a b) (.>.) :: (Columnable a, Ord a) => TExpr cols a -> TExpr cols a -> TExpr cols Bool -(.>.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (>) "gt" (Just ">") False 4) a b) +(.>.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (>) (Builtin BGt) (Just ">") False 4) a b) -- Same-type arithmetic operators @@ -306,18 +306,18 @@ infixr 8 .^^., .^^, .^., .^ (.^^.) :: (Columnable a, Columnable b, Fractional a, Integral b) => TExpr cols a -> TExpr cols b -> TExpr cols a -(.^^.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (^^) "pow" (Just ".^^.") False 8) a b) +(.^^.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (^^) (Builtin BPowFrac) (Just ".^^.") False 8) a b) (.^.) :: (Columnable a, Columnable b, Num a, Integral b) => TExpr cols a -> TExpr cols b -> TExpr cols a -(.^.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (^) "pow" (Just ".^.") False 8) a b) +(.^.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (^) (Builtin BPow) (Just ".^.") False 8) a b) (.&&.) :: TExpr cols Bool -> TExpr cols Bool -> TExpr cols Bool -(.&&.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (&&) "and" (Just ".&&.") True 3) a b) +(.&&.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (&&) (Builtin BAnd) (Just ".&&.") True 3) a b) (.||.) :: TExpr cols Bool -> TExpr cols Bool -> TExpr cols Bool -(.||.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (||) "or" (Just ".||.") True 2) a b) +(.||.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (||) (Builtin BOr) (Just ".||.") True 2) a b) -- | Nullable-aware logical AND; @Maybe Bool@ when either operand is nullable. (.&&) :: @@ -326,7 +326,8 @@ infixr 8 .^^., .^^, .^., .^ TExpr cols b -> TExpr cols (NullCmpResult a b) (.&&) (TExpr a) (TExpr b) = - TExpr (Binary (MkBinaryOp (nullCmpOp (&&)) "nulland" (Just ".&&") True 3) a b) + TExpr + (Binary (MkBinaryOp (nullCmpOp (&&)) (Builtin BNullAnd) (Just ".&&") True 3) a b) -- | Nullable-aware logical OR; @Maybe Bool@ when either operand is nullable. (.||) :: @@ -335,7 +336,8 @@ infixr 8 .^^., .^^, .^., .^ TExpr cols b -> TExpr cols (NullCmpResult a b) (.||) (TExpr a) (TExpr b) = - TExpr (Binary (MkBinaryOp (nullCmpOp (||)) "nullor" (Just ".||") True 2) a b) + TExpr + (Binary (MkBinaryOp (nullCmpOp (||)) (Builtin BNullOr) (Just ".||") True 2) a b) infixl 6 .+, .- infixl 7 .*, ./ @@ -352,7 +354,7 @@ infixl 7 .*, ./ (.+) (TExpr a) (TExpr b) = TExpr ( Binary - (MkBinaryOp (applyNull2 (widenArithOp (+))) "nulladd" (Just "+") True 6) + (MkBinaryOp (applyNull2 (widenArithOp (+))) (Builtin BNullAdd) (Just "+") True 6) a b ) @@ -369,7 +371,7 @@ infixl 7 .*, ./ (.-) (TExpr a) (TExpr b) = TExpr ( Binary - (MkBinaryOp (applyNull2 (widenArithOp (-))) "nullsub" (Just "-") False 6) + (MkBinaryOp (applyNull2 (widenArithOp (-))) (Builtin BNullSub) (Just "-") False 6) a b ) @@ -386,7 +388,7 @@ infixl 7 .*, ./ (.*) (TExpr a) (TExpr b) = TExpr ( Binary - (MkBinaryOp (applyNull2 (widenArithOp (*))) "nullmul" (Just "*") True 7) + (MkBinaryOp (applyNull2 (widenArithOp (*))) (Builtin BNullMul) (Just "*") True 7) a b ) @@ -403,7 +405,7 @@ infixl 7 .*, ./ (./) (TExpr a) (TExpr b) = TExpr ( Binary - (MkBinaryOp (applyNull2 (divArithOp (/))) "nulldiv" (Just "/") False 7) + (MkBinaryOp (applyNull2 (divArithOp (/))) (Builtin BNullDiv) (Just "/") False 7) a b ) @@ -420,7 +422,12 @@ infixl 7 .*, ./ ) => TExpr cols a -> TExpr cols b -> TExpr cols a (.^^) (TExpr a) (TExpr b) = - TExpr (Binary (MkBinaryOp (applyNull2 (^^)) "pow" (Just ".^^") False 8) a b) + TExpr + ( Binary + (MkBinaryOp (applyNull2 (^^)) (Builtin BPowFrac) (Just ".^^") False 8) + a + b + ) -- | Nullable-aware exponentiation (num base, integral exponent). (.^) :: @@ -434,7 +441,8 @@ infixl 7 .*, ./ ) => TExpr cols a -> TExpr cols b -> TExpr cols a (.^) (TExpr a) (TExpr b) = - TExpr (Binary (MkBinaryOp (applyNull2 (^)) "pow" (Just ".^") False 8) a b) + TExpr + (Binary (MkBinaryOp (applyNull2 (^)) (Builtin BPow) (Just ".^") False 8) a b) -- | Nullable-aware equality that widens numeric operands to a common type. (.==) :: @@ -447,7 +455,11 @@ infixl 7 .*, ./ TExpr cols (NullCmpResult a b) (.==) (TExpr a) (TExpr b) = TExpr - (Binary (MkBinaryOp (applyNull2 (widenCmpOp (==))) "eq" (Just "==") True 4) a b) + ( Binary + (MkBinaryOp (applyNull2 (widenCmpOp (==))) (Builtin BEq) (Just "==") True 4) + a + b + ) -- | Nullable-aware inequality. Widens numeric operands to their common type. (./=) :: @@ -460,7 +472,11 @@ infixl 7 .*, ./ TExpr cols (NullCmpResult a b) (./=) (TExpr a) (TExpr b) = TExpr - (Binary (MkBinaryOp (applyNull2 (widenCmpOp (/=))) "neq" (Just "/=") True 4) a b) + ( Binary + (MkBinaryOp (applyNull2 (widenCmpOp (/=))) (Builtin BNeq) (Just "/=") True 4) + a + b + ) -- | Nullable-aware less-than. Widens numeric operands to their common type. (.<) :: @@ -473,7 +489,11 @@ infixl 7 .*, ./ TExpr cols (NullCmpResult a b) (.<) (TExpr a) (TExpr b) = TExpr - (Binary (MkBinaryOp (applyNull2 (widenCmpOp (<))) "lt" (Just "<") False 4) a b) + ( Binary + (MkBinaryOp (applyNull2 (widenCmpOp (<))) (Builtin BLt) (Just "<") False 4) + a + b + ) -- | Nullable-aware <=. Widens numeric operands to their common type. (.<=) :: @@ -486,7 +506,11 @@ infixl 7 .*, ./ TExpr cols (NullCmpResult a b) (.<=) (TExpr a) (TExpr b) = TExpr - (Binary (MkBinaryOp (applyNull2 (widenCmpOp (<=))) "leq" (Just "<=") False 4) a b) + ( Binary + (MkBinaryOp (applyNull2 (widenCmpOp (<=))) (Builtin BLeq) (Just "<=") False 4) + a + b + ) -- | Nullable-aware >=. Widens numeric operands to their common type. (.>=) :: @@ -499,7 +523,11 @@ infixl 7 .*, ./ TExpr cols (NullCmpResult a b) (.>=) (TExpr a) (TExpr b) = TExpr - (Binary (MkBinaryOp (applyNull2 (widenCmpOp (>=))) "geq" (Just ">=") False 4) a b) + ( Binary + (MkBinaryOp (applyNull2 (widenCmpOp (>=))) (Builtin BGeq) (Just ">=") False 4) + a + b + ) -- | Nullable-aware greater-than. Widens numeric operands to their common type. (.>) :: @@ -512,10 +540,14 @@ infixl 7 .*, ./ TExpr cols (NullCmpResult a b) (.>) (TExpr a) (TExpr b) = TExpr - (Binary (MkBinaryOp (applyNull2 (widenCmpOp (>))) "gt" (Just ">") False 4) a b) + ( Binary + (MkBinaryOp (applyNull2 (widenCmpOp (>))) (Builtin BGt) (Just ">") False 4) + a + b + ) not :: TExpr cols Bool -> TExpr cols Bool -not (TExpr e) = TExpr (Unary (MkUnaryOp Prelude.not "not" (Just "!")) e) +not (TExpr e) = TExpr (F.not e) sum :: (Columnable a, Num a) => TExpr cols a -> TExpr cols a sum (TExpr e) = TExpr (F.sum e) @@ -552,43 +584,23 @@ over (TExpr e) = TExpr{unTExpr = F.over (symbolVals @names) e} castExpr :: forall b cols src. (Columnable b, Columnable src, Read b) => TExpr cols src -> TExpr cols (Maybe b) -castExpr (TExpr e) = - TExpr - (CastExprWith @b @(Maybe b) @src "castExpr" (either (const Nothing) Just) e) +castExpr (TExpr e) = TExpr (F.castExpr @b e) castExprWithDefault :: forall b cols src. (Columnable b, Columnable src, Read b) => b -> TExpr cols src -> TExpr cols b -castExprWithDefault def (TExpr e) = - TExpr - ( CastExprWith @b @b @src - ("castExprWithDefault:" <> T.pack (show def)) - (fromRight def) - e - ) +castExprWithDefault def (TExpr e) = TExpr (F.castExprWithDefault def e) castExprEither :: forall b cols src. (Columnable b, Columnable src, Read b) => TExpr cols src -> TExpr cols (Either T.Text b) -castExprEither (TExpr e) = - TExpr - ( CastExprWith @b @(Either T.Text b) @src - "castExprEither" - (either (Left . T.pack) Right) - e - ) +castExprEither (TExpr e) = TExpr (F.castExprEither @b e) unsafeCastExpr :: forall b cols src. (Columnable b, Columnable src, Read b) => TExpr cols src -> TExpr cols b -unsafeCastExpr (TExpr e) = - TExpr - ( CastExprWith @b @b @src - "unsafeCastExpr" - (fromRight (error "unsafeCastExpr: unexpected Nothing in column")) - e - ) +unsafeCastExpr (TExpr e) = TExpr (F.unsafeCastExpr @b e) toDouble :: (Columnable a, Real a) => TExpr cols a -> TExpr cols Double toDouble (TExpr e) = TExpr (F.toDouble e) diff --git a/dataframe-operations/src/DataFrame/Typed/Expr/Extra.hs b/dataframe-operations/src/DataFrame/Typed/Expr/Extra.hs index 0d46f0da..bcb2244f 100644 --- a/dataframe-operations/src/DataFrame/Typed/Expr/Extra.hs +++ b/dataframe-operations/src/DataFrame/Typed/Expr/Extra.hs @@ -1,9 +1,33 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE DataKinds #-} {-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE KindSignatures #-} {-# LANGUAGE MonoLocalBinds #-} {-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} -- | Typed counterparts of the remaining "DataFrame.Functions" combinators. module DataFrame.Typed.Expr.Extra ( + cast, + castWithDefault, + castEither, + unsafeCast, + eq, + lt, + gt, + leq, + geq, + and, + or, + countDistinct, + topK, + bottomK, + isNull, + isNotNull, + contains, + length, div, mod, mode, @@ -40,15 +64,116 @@ module DataFrame.Typed.Expr.Extra ( bind, ) where +import Data.Proxy (Proxy (..)) import qualified Data.Text as T import Data.Time (Day, ParseTime) import qualified Data.Vector.Unboxed as VU -import Prelude hiding (div, max, min, mod) +import GHC.TypeLits (KnownSymbol, Symbol, symbolVal) +import Prelude hiding (and, div, length, max, min, mod, or) import qualified DataFrame.Functions as F import DataFrame.Internal.Column (Columnable) +import DataFrame.Internal.Expression.Operators.Nullable (BaseType) +import DataFrame.Typed.Schema (AssertPresent) import DataFrame.Typed.Types (TExpr (..)) +colName :: forall (name :: Symbol). (KnownSymbol name) => T.Text +colName = T.pack (symbolVal (Proxy @name)) + +-- | Parse a schema-checked column as @a@; failures become 'Nothing'. +cast :: + forall (name :: Symbol) a cols. + (KnownSymbol name, AssertPresent name cols, Columnable a, Read a) => + TExpr cols (Maybe a) +cast = TExpr (F.cast @a (colName @name)) + +-- | Parse a schema-checked column as @a@, with a default for failures. +castWithDefault :: + forall (name :: Symbol) a cols. + (KnownSymbol name, AssertPresent name cols, Columnable a, Read a) => + a -> TExpr cols a +castWithDefault d = TExpr (F.castWithDefault d (colName @name)) + +-- | Parse a schema-checked column as @a@, keeping failures as 'Left'. +castEither :: + forall (name :: Symbol) a cols. + (KnownSymbol name, AssertPresent name cols, Columnable a, Read a) => + TExpr cols (Either T.Text a) +castEither = TExpr (F.castEither @a (colName @name)) + +-- | Parse a schema-checked column as @a@, erroring on failures. +unsafeCast :: + forall (name :: Symbol) a cols. + (KnownSymbol name, AssertPresent name cols, Columnable a, Read a) => + TExpr cols a +unsafeCast = TExpr (F.unsafeCast @a (colName @name)) + +-- | Equality. +eq :: + (Columnable a, Eq a, a ~ BaseType a) => + TExpr cols a -> TExpr cols a -> TExpr cols Bool +eq (TExpr a) (TExpr b) = TExpr (F.eq a b) + +-- | Less than. +lt :: + (Columnable a, Ord a, a ~ BaseType a) => + TExpr cols a -> TExpr cols a -> TExpr cols Bool +lt (TExpr a) (TExpr b) = TExpr (F.lt a b) + +-- | Greater than. +gt :: + (Columnable a, Ord a, a ~ BaseType a) => + TExpr cols a -> TExpr cols a -> TExpr cols Bool +gt (TExpr a) (TExpr b) = TExpr (F.gt a b) + +-- | Less than or equal. +leq :: + (Columnable a, Ord a, a ~ BaseType a) => + TExpr cols a -> TExpr cols a -> TExpr cols Bool +leq (TExpr a) (TExpr b) = TExpr (F.leq a b) + +-- | Greater than or equal. +geq :: + (Columnable a, Ord a, a ~ BaseType a) => + TExpr cols a -> TExpr cols a -> TExpr cols Bool +geq (TExpr a) (TExpr b) = TExpr (F.geq a b) + +-- | Logical and. +and :: TExpr cols Bool -> TExpr cols Bool -> TExpr cols Bool +and (TExpr a) (TExpr b) = TExpr (F.and a b) + +-- | Logical or. +or :: TExpr cols Bool -> TExpr cols Bool -> TExpr cols Bool +or (TExpr a) (TExpr b) = TExpr (F.or a b) + +-- | Number of distinct values (aggregation). +countDistinct :: (Columnable a, Ord a) => TExpr cols a -> TExpr cols Int +countDistinct (TExpr e) = TExpr (F.countDistinct e) + +-- | The k largest values per group (aggregation). +topK :: (Columnable a, Ord a) => Int -> TExpr cols a -> TExpr cols [a] +topK k (TExpr e) = TExpr (F.topK k e) + +-- | The k smallest values per group (aggregation). +bottomK :: (Columnable a, Ord a) => Int -> TExpr cols a -> TExpr cols [a] +bottomK k (TExpr e) = TExpr (F.bottomK k e) + +-- | SQL spelling of 'isNothing'. +isNull :: (Columnable a) => TExpr cols (Maybe a) -> TExpr cols Bool +isNull (TExpr e) = TExpr (F.isNull e) + +-- | SQL spelling of 'isJust'. +isNotNull :: (Columnable a) => TExpr cols (Maybe a) -> TExpr cols Bool +isNotNull (TExpr e) = TExpr (F.isNotNull e) + +-- | Whether the text contains the literal substring. +contains :: T.Text -> TExpr cols T.Text -> TExpr cols Bool +contains needle (TExpr e) = TExpr (F.contains needle e) + +-- | Character count. +length :: TExpr cols T.Text -> TExpr cols Int +length (TExpr e) = TExpr (F.length e) + -- | Integer division. div :: (Integral a, Columnable a) => TExpr cols a -> TExpr cols a -> TExpr cols a diff --git a/dataframe-parquet/src/DataFrame/IO/Parquet.hs b/dataframe-parquet/src/DataFrame/IO/Parquet.hs index 2bac302e..1a67a37a 100644 --- a/dataframe-parquet/src/DataFrame/IO/Parquet.hs +++ b/dataframe-parquet/src/DataFrame/IO/Parquet.hs @@ -175,8 +175,10 @@ import DataFrame.Internal.Column.Builder ( import DataFrame.Internal.Control.Concurrent (forkJoin_) import DataFrame.Internal.DataFrame (DataFrame (..)) import DataFrame.Internal.Expression ( - BinaryOp (binaryName), + BinOpId (..), + BinaryOp (binaryId), Expr (..), + builtinOf, getColumns, ) import DataFrame.Operations.Merge () @@ -419,7 +421,8 @@ rowGroupMayMatch metadata rg = go chunks = unField (rg_columns rg) go :: Expr Bool -> Bool go (Binary op l r) - | binaryName op == "and" = case (l, r) of + | Just tag <- builtinOf (binaryId op) + , tag `elem` [BAnd, BNullAnd] = case (l, r) of (l' :: Expr lt, r' :: Expr rt) | Just Refl <- testEquality (typeRep @lt) (typeRep @Bool) , Just Refl <- testEquality (typeRep @rt) (typeRep @Bool) -> @@ -429,31 +432,31 @@ rowGroupMayMatch metadata rg = go (Col name :: Expr lt, Lit (v :: rt)) | Just Refl <- testEquality (typeRep @lt) (typeRep @Int) , Just Refl <- testEquality (typeRep @rt) (typeRep @Int) -> - decide name (binaryName op) v + maybe True (\t -> decide name t v) (builtinOf (binaryId op)) (Lit (v :: lt), Col name :: Expr rt) | Just Refl <- testEquality (typeRep @lt) (typeRep @Int) , Just Refl <- testEquality (typeRep @rt) (typeRep @Int) -> - decide name (flipped (binaryName op)) v + maybe True (\t -> decide name (flipped t) v) (builtinOf (binaryId op)) _ -> True go _ = True - flipped :: T.Text -> T.Text + flipped :: BinOpId -> BinOpId flipped n = case n of - "lt" -> "gt" - "gt" -> "lt" - "leq" -> "geq" - "geq" -> "leq" + BLt -> BGt + BGt -> BLt + BLeq -> BGeq + BGeq -> BLeq other -> other -- @column v@ can hold for some row in [lo, hi]. - decide :: T.Text -> T.Text -> Int -> Bool + decide :: T.Text -> BinOpId -> Int -> Bool decide name cmp v = case columnRange name of Nothing -> True Just (lo, hi) -> case cmp of - "eq" -> lo <= v && v <= hi - "neq" -> not (lo == v && hi == v) - "lt" -> lo < v - "leq" -> lo <= v - "gt" -> hi > v - "geq" -> hi >= v + BEq -> lo <= v && v <= hi + BNeq -> not (lo == v && hi == v) + BLt -> lo < v + BLeq -> lo <= v + BGt -> hi > v + BGeq -> hi >= v _ -> True columnRange :: T.Text -> Maybe (Int, Int) columnRange name = do diff --git a/dataframe-th/dataframe-th.cabal b/dataframe-th/dataframe-th.cabal index d5069da5..860ae9ea 100644 --- a/dataframe-th/dataframe-th.cabal +++ b/dataframe-th/dataframe-th.cabal @@ -27,6 +27,7 @@ common warnings library import: warnings exposed-modules: + DataFrame.TH.Identifier DataFrame.TH.Records DataFrame.Typed.TH.Records -- Internal record→schema TH; sealed (superseded publicly by diff --git a/dataframe-th/src/DataFrame/TH/Identifier.hs b/dataframe-th/src/DataFrame/TH/Identifier.hs new file mode 100644 index 00000000..d0d7cbbd --- /dev/null +++ b/dataframe-th/src/DataFrame/TH/Identifier.hs @@ -0,0 +1,75 @@ +{-# LANGUAGE OverloadedStrings #-} + +-- | Turning column names into valid Haskell identifiers for generated code. +module DataFrame.TH.Identifier ( + sanitize, + isHaskellIdentifier, + isReservedId, + isVarId, +) where + +import qualified Data.Char as Char +import qualified Data.Text as T + +-- | Haskell Report 2.4 reserved words. +isReservedId :: T.Text -> Bool +isReservedId t = case t of + "case" -> True + "class" -> True + "data" -> True + "default" -> True + "deriving" -> True + "do" -> True + "else" -> True + "foreign" -> True + "if" -> True + "import" -> True + "in" -> True + "infix" -> True + "infixl" -> True + "infixr" -> True + "instance" -> True + "let" -> True + "module" -> True + "newtype" -> True + "of" -> True + "then" -> True + "type" -> True + "where" -> True + _ -> False + +isVarId :: T.Text -> Bool +isVarId t = case T.uncons t of + -- '_' is skipped on purpose so "Name" and "_Name_" cannot collide. + Just (c, _) -> Char.isLower c && Char.isAlpha c + Nothing -> False + +isHaskellIdentifier :: T.Text -> Bool +isHaskellIdentifier t = Prelude.not (isVarId t) || isReservedId t + +sanitize :: T.Text -> T.Text +sanitize t + | isValid = t + | isHaskellIdentifier t' = "_" <> t' <> "_" + | otherwise = t' + where + isValid = + Prelude.not (isHaskellIdentifier t) + && isVarId t + && T.all Char.isAlphaNum t + t' = T.map replaceInvalidCharacters . T.filter (Prelude.not . parentheses) $ t + replaceInvalidCharacters c + | Char.isUpper c = Char.toLower c + | Char.isSpace c = '_' + | Char.isPunctuation c = '_' -- '-' will also become a '_' + | Char.isSymbol c = '_' + | Char.isAlphaNum c = c -- Blanket condition + | otherwise = '_' -- If we're unsure we'll default to an underscore + parentheses c = case c of + '(' -> True + ')' -> True + '{' -> True + '}' -> True + '[' -> True + ']' -> True + _ -> False diff --git a/dataframe-th/src/DataFrame/TH/Records.hs b/dataframe-th/src/DataFrame/TH/Records.hs index 6124a097..c9b01b23 100644 --- a/dataframe-th/src/DataFrame/TH/Records.hs +++ b/dataframe-th/src/DataFrame/TH/Records.hs @@ -41,7 +41,6 @@ import Language.Haskell.TH import qualified Language.Haskell.TH.Syntax as TH import DataFrame.Expression.Operators (col) -import DataFrame.Functions (sanitize) import DataFrame.Internal.Column (columnTypeString) import DataFrame.Internal.DataFrame ( DataFrame (..), @@ -49,6 +48,7 @@ import DataFrame.Internal.DataFrame ( ) import DataFrame.Internal.Expression (Expr) import DataFrame.Internal.Schema.TH (schemaTypeE) +import DataFrame.TH.Identifier (sanitize) import DataFrame.Schema (Schema, makeSchema) import Prelude as P diff --git a/dataframe-viz/src/DataFrame/Display/Internal/VegaLite.hs b/dataframe-viz/src/DataFrame/Display/Internal/VegaLite.hs index 4195ba67..3312c405 100644 --- a/dataframe-viz/src/DataFrame/Display/Internal/VegaLite.hs +++ b/dataframe-viz/src/DataFrame/Display/Internal/VegaLite.hs @@ -229,9 +229,7 @@ materialiseExpr df expr = either throw id $ do v <- eval (FlatCtx df) expr pure $ materialize @a (fst (dataframeDimensions df)) v -{- | A column's values as JSON, with null slots as 'Null'. Vega-Lite groups a -null category as its own bar and drops null positions from a numeric scale. --} +-- | A column's values as JSON, with null slots as 'Null'. columnToValues :: FieldType -> Column -> [Value] columnToValues ft col = case col of MergedColumn _ _ -> columnToValues ft (materializeMerged col) @@ -329,10 +327,7 @@ transformValue (RegressionT yField xField) = transformValue (DensityT field) = object ["density" .= fieldRef field] -{- | A column name as a Vega-Lite field reference. Vega-Lite reads @.@, @[@ and -@]@ in a field string as nested access, so a column called @sepal.length@ -must be written @sepal\\.length@ or every datum resolves to undefined. --} +-- | A column as a Vega-Lite field; escapes @.@, @[@, @]@ (else nested). fieldRef :: T.Text -> T.Text fieldRef = T.concatMap esc where diff --git a/dataframe-viz/src/DataFrame/Display/Web/Chart/Typed.hs b/dataframe-viz/src/DataFrame/Display/Web/Chart/Typed.hs index 596b05c4..71fabd51 100644 --- a/dataframe-viz/src/DataFrame/Display/Web/Chart/Typed.hs +++ b/dataframe-viz/src/DataFrame/Display/Web/Chart/Typed.hs @@ -63,10 +63,7 @@ import DataFrame.Internal.Expression.Operators.Nullable (BaseType) import DataFrame.Typed.Types (TExpr (..), TypedDataFrame (..)) import GHC.TypeLits (Symbol) -{- | A typed chart: a phantom-@cols@ wrapper over the untyped builder. -Encoded expressions must be non-nullable (@a ~ BaseType a@); drop nulls first -with 'DataFrame.Typed.filterJust' or fill them with 'DataFrame.Typed.impute'. --} +-- | A typed chart over the untyped builder; encoded exprs must be non-null. newtype Chart (cols :: [(Symbol, Type)]) = Chart C.Chart -- | Start a typed chart from a typed frame. diff --git a/dataframe.cabal b/dataframe.cabal index b4bc2212..bc44f6a1 100644 --- a/dataframe.cabal +++ b/dataframe.cabal @@ -319,6 +319,7 @@ test-suite tests Operations.Inference, Operations.InsertColumn, Operations.Join, + Operations.OperatorIdentity, Operations.Merge, Operations.Nullable, Operations.NullableHashing, @@ -365,6 +366,7 @@ test-suite tests dataframe-operations >= 2.6 && < 2.7, dataframe-parquet >= 1.6 && < 1.7, dataframe-parsing >= 2.2 && < 2.3, + dataframe-th >= 2.3 && < 2.4, HUnit >= 1.6 && < 1.8, QuickCheck >= 2 && < 3, random-shuffle >= 0.0.4 && < 1, diff --git a/examples/examples.cabal b/examples/examples.cabal index 06558b35..eb30f708 100644 --- a/examples/examples.cabal +++ b/examples/examples.cabal @@ -90,6 +90,7 @@ executable examples DataFrame.Internal.Binary, DataFrame.Internal.Expression.Operators.Nullable, DataFrame.Internal.Expression.Operators, + DataFrame.Internal.Expression.OpId, DataFrame.Expression.Operators, DataFrame.DecisionTree, DataFrame.Typed.Types, @@ -111,6 +112,7 @@ executable examples DataFrame.Internal.Schema.TH, DataFrame.TH.CSV, DataFrame.TH.Parquet, + DataFrame.TH.Identifier, DataFrame.TH.Records, DataFrame.Typed.TH.CSV, DataFrame.Typed.TH.Records, diff --git a/src/DataFrame/Typed.hs b/src/DataFrame/Typed.hs index 2f05dd68..09f1d050 100644 --- a/src/DataFrame/Typed.hs +++ b/src/DataFrame/Typed.hs @@ -59,6 +59,20 @@ module DataFrame.Typed ( over, -- * Expression combinators (full DataFrame.Functions parity) + DataFrame.Typed.Expr.eq, + DataFrame.Typed.Expr.lt, + DataFrame.Typed.Expr.gt, + DataFrame.Typed.Expr.leq, + DataFrame.Typed.Expr.geq, + DataFrame.Typed.Expr.and, + DataFrame.Typed.Expr.or, + DataFrame.Typed.Expr.countDistinct, + DataFrame.Typed.Expr.topK, + DataFrame.Typed.Expr.bottomK, + DataFrame.Typed.Expr.isNull, + DataFrame.Typed.Expr.isNotNull, + DataFrame.Typed.Expr.contains, + DataFrame.Typed.Expr.length, DataFrame.Typed.Expr.div, DataFrame.Typed.Expr.mod, mode, @@ -95,6 +109,10 @@ module DataFrame.Typed ( bind, -- * Cast / coercion expressions + DataFrame.Typed.Expr.cast, + DataFrame.Typed.Expr.castWithDefault, + DataFrame.Typed.Expr.castEither, + DataFrame.Typed.Expr.unsafeCast, castExpr, castExprWithDefault, castExprEither, diff --git a/tests/Functions.hs b/tests/Functions.hs index c156cc8d..0f5f334e 100644 --- a/tests/Functions.hs +++ b/tests/Functions.hs @@ -5,11 +5,9 @@ module Functions where import Data.Time.Calendar (Day, fromGregorian) import qualified DataFrame as D -import DataFrame.Functions ( - sanitize, - ) import qualified DataFrame.Functions as F import qualified DataFrame.Internal.Column as DI +import DataFrame.TH.Identifier (sanitize) import Test.HUnit -- Test cases for the sanitize function diff --git a/tests/IR/ExprJsonRoundtrip.hs b/tests/IR/ExprJsonRoundtrip.hs index b31b9895..c2682c46 100644 --- a/tests/IR/ExprJsonRoundtrip.hs +++ b/tests/IR/ExprJsonRoundtrip.hs @@ -82,6 +82,7 @@ exprRoundtripTests = , testCase "division" (roundtrips (F.col @Double "x" / F.lit 2)) , testCase "comparison" (roundtrips (F.col @Double "x" .>. F.lit 0)) , testCase "unary sqrt" (roundtrips (sqrt (F.col @Double "x"))) + , testCase "logBase" (roundtrips (logBase (F.lit 2) (F.col @Double "x"))) , testCase "if/then/else" ( roundtrips @@ -101,6 +102,8 @@ aggRoundtripTests = , testCase "mean" (roundtrips (F.mean (F.col @Double "x"))) , testCase "variance" (roundtrips (F.variance (F.col @Double "x"))) , testCase "median" (roundtrips (F.median (F.col @Double "x"))) + , testCase "stddev" (roundtrips (F.stddev (F.col @Double "x"))) + , testCase "countDistinct" (roundtrips (F.countDistinct (F.col @Int "x"))) , testCase "over" (roundtrips (F.over ["g"] (F.sum (F.col @Double "x")))) ] @@ -110,6 +113,9 @@ limitationTests = [ testCase "collect cannot encode" (isLeft (encodeExpr (F.collect (F.col @Double "x")))) + , testCase + "anonymous udf cannot encode" + (isLeft (encodeExpr (F.lift (+ 1) (F.col @Double "x")))) ] -- Pin the wire shape Python decodes against: (x * 2) + 1 over doubles. diff --git a/tests/Main.hs b/tests/Main.hs index 802c29fc..93ede77f 100644 --- a/tests/Main.hs +++ b/tests/Main.hs @@ -47,6 +47,7 @@ import qualified Operations.Merge import qualified Operations.MonomorphicLoops import qualified Operations.Nullable import qualified Operations.NullableHashing +import qualified Operations.OperatorIdentity import qualified Operations.ParallelGroupBy import qualified Operations.ParallelJoin import qualified Operations.Provenance @@ -109,6 +110,7 @@ tests = ++ Operations.Inference.tests ++ Operations.InsertColumn.tests ++ Operations.Join.tests + ++ Operations.OperatorIdentity.tests ++ Operations.Merge.tests ++ Operations.Nullable.tests ++ Operations.NullableHashing.tests diff --git a/tests/Operations/OperatorIdentity.hs b/tests/Operations/OperatorIdentity.hs new file mode 100644 index 00000000..625eb689 --- /dev/null +++ b/tests/Operations/OperatorIdentity.hs @@ -0,0 +1,255 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +-- | Operators are identified by tag, never by display name. +module Operations.OperatorIdentity (tests) where + +import qualified Data.Text as T +import Data.Time (Day) +import qualified DataFrame as D +import DataFrame.Expression.Operators +import qualified DataFrame.Functions as F +import DataFrame.Internal.Column (Columnable) +import DataFrame.Internal.Expression (Expr, eqExpr, structuralKey) +import DataFrame.Internal.Expression.Simplify (simplify) + +import Test.HUnit + +xs :: D.DataFrame +xs = D.fromColumns [("x", D.fromList [1, 2, 3 :: Int])] + +derived :: Expr Int -> [Int] +derived e = D.columnAsList (F.col @Int "y") (D.derive "y" e xs) + +weird :: T.Text -> Expr Int -> Expr Int -> Expr Int +weird opName = lift2Decorated (\a b -> a * 100 + b) opName Nothing True 6 + +udfNamedLikeBuiltin :: Test +udfNamedLikeBuiltin = + TestLabel "udf named add keeps its own function" . TestCase $ + assertEqual "" [101, 202, 303] (derived (weird "add" x x)) + where + x = F.col @Int "x" + +udfOtherName :: Test +udfOtherName = + TestLabel "udf with fresh name" . TestCase $ + assertEqual "" [101, 202, 303] (derived (weird "myadd" x x)) + where + x = F.col @Int "x" + +anonymousNotEqual :: Test +anonymousNotEqual = + TestLabel "anonymous udfs are not equal" . TestCase $ + assertBool "" (not (eqExpr (F.lift2 (+) x x) (F.lift2 (*) x x))) + where + x = F.col @Int "x" + +namedEqual :: Test +namedEqual = + TestLabel "same named op is equal" . TestCase $ + assertBool "" (eqExpr (weird "f" x x) (weird "f" x x)) + where + x = F.col @Int "x" + +namedVsBuiltin :: Test +namedVsBuiltin = + TestLabel "named add differs from builtin add" . TestCase $ + assertBool "" (not (eqExpr (weird "add" x x) (x + x))) + where + x = F.col @Int "x" + +fromMaybeDefaults :: Test +fromMaybeDefaults = + TestLabel "fromMaybe defaults are part of identity" . TestCase $ + assertBool "" (not (eqExpr (F.fromMaybe 0 m) (F.fromMaybe 1 m))) + where + m = F.col @(Maybe Int) "m" + +powVariants :: Test +powVariants = + TestLabel "^ and ^^ are different ops" . TestCase $ + assertBool "" (not (eqExpr (d .^. two) (d .^^. two))) + where + d = F.col @Double "d" + two = F.lit (2 :: Int) + +rowsKept :: D.DataFrame -> Expr Bool -> Int +rowsKept df p = D.nRows (D.filterWhere p df) + +-- Both sides must keep the same rows before and after simplification. +sameRows :: String -> D.DataFrame -> Expr Bool -> Test +sameRows label df p = + TestLabel label . TestCase $ + assertEqual label (rowsKept df p) (rowsKept df (simplify p)) + +anonymousAtomsNotMerged :: Test +anonymousAtomsNotMerged = + sameRows "distinct udfs over one column are separate atoms" ds $ + F.and + (F.lift (* 10) d .> F.lit (3 :: Double)) + (F.lift id d .< F.lit (2 :: Double)) + where + ds = D.fromColumns [("d", D.fromList [1, 2, 3 :: Double])] + d = F.col @Double "d" + +fromMaybeTrueNotFalseOnNull :: Test +fromMaybeTrueNotFalseOnNull = + sameRows "fromMaybe True keeps null rows" ms $ + F.and + (F.fromMaybe True (m .< F.lit (5 :: Double))) + (F.fromMaybe True (m .> F.lit (10 :: Double))) + where + ms = + D.fromColumns + [("m", D.fromList [Just 1, Nothing, Just 20 :: Maybe Double])] + m = F.col @(Maybe Double) "m" + +-- The six-sum moment shape must not see through a user function. +momentsRespectUdf :: Test +momentsRespectUdf = + TestLabel "moment fast path keeps udf in sum" . TestCase $ + assertEqual "" [60, 150] (D.columnAsList (F.col @Double "sx") out) + where + df = + D.fromColumns + [ ("k", D.fromList [1, 1, 1, 2, 2, 2 :: Int]) + , ("x", D.fromList [1, 2, 3, 4, 5, 6 :: Double]) + , ("y", D.fromList [6, 5, 4, 3, 2, 1 :: Double]) + ] + x = F.col @Double "x" + y = F.col @Double "y" + out = + D.sortBy [D.Asc (F.col @Int "k")] $ + D.aggregate + [ F.count x `as` "n" + , F.sum (F.lift (* 10) x) `as` "sx" + , F.sum y `as` "sy" + , F.sum (x * x) `as` "sxx" + , F.sum (y * y) `as` "syy" + , F.sum (x * y) `as` "sxy" + ] + (D.groupBy ["k"] df) + +-- A function is builtin iff its expression has a provable structural key. +hasKey :: (Columnable a) => Expr a -> Bool +hasKey = (/= Nothing) . structuralKey + +-- Every non-higher-order function in DataFrame.Functions is a builtin. +everyFunctionIsBuiltin :: Test +everyFunctionIsBuiltin = + TestList + [ TestLabel fn (TestCase (assertBool fn ok)) + | (fn, ok) <- builtins + ] + where + x = F.col @Int "x" + d = F.col @Double "d" + m = F.col @(Maybe Int) "m" + t = F.col @T.Text "t" + b = F.col @Bool "b" + l = F.col @[Int] "l" + day = F.col @Day "day" + builtins = + [ ("cast", hasKey (F.cast @Int "x")) + , ("castWithDefault", hasKey (F.castWithDefault (0 :: Int) "x")) + , ("castEither", hasKey (F.castEither @Int "x")) + , ("unsafeCast", hasKey (F.unsafeCast @Int "x")) + , ("castExpr", hasKey (F.castExpr @Int t)) + , ("castExprWithDefault", hasKey (F.castExprWithDefault (0 :: Int) t)) + , ("castExprEither", hasKey (F.castExprEither @Int t)) + , ("unsafeCastExpr", hasKey (F.unsafeCastExpr @Int t)) + , ("toDouble", hasKey (F.toDouble x)) + , ("div", hasKey (F.div x x)) + , ("mod", hasKey (F.mod x x)) + , ("eq", hasKey (F.eq x x)) + , ("lt", hasKey (F.lt x x)) + , ("gt", hasKey (F.gt x x)) + , ("leq", hasKey (F.leq x x)) + , ("geq", hasKey (F.geq x x)) + , ("and", hasKey (F.and b b)) + , ("or", hasKey (F.or b b)) + , ("not", hasKey (F.not b)) + , ("count", hasKey (F.count x)) + , ("countDistinct", hasKey (F.countDistinct x)) + , ("countAll", hasKey F.countAll) + , ("collect", hasKey (F.collect x)) + , ("mode", hasKey (F.mode x)) + , ("minimum", hasKey (F.minimum x)) + , ("maximum", hasKey (F.maximum x)) + , ("sum", hasKey (F.sum x)) + , ("sumMaybe", hasKey (F.sumMaybe m)) + , ("mean", hasKey (F.mean x)) + , ("topK", hasKey (F.topK 2 x)) + , ("bottomK", hasKey (F.bottomK 2 x)) + , ("meanMaybe", hasKey (F.meanMaybe m)) + , ("variance", hasKey (F.variance x)) + , ("median", hasKey (F.median x)) + , ("medianMaybe", hasKey (F.medianMaybe m)) + , ("percentile", hasKey (F.percentile 50 d)) + , ("stddev", hasKey (F.stddev x)) + , ("stddevMaybe", hasKey (F.stddevMaybe m)) + , ("zScore", hasKey (F.zScore d)) + , ("pow", hasKey (F.pow x 2)) + , ("relu", hasKey (F.relu x)) + , ("min", hasKey (F.min x x)) + , ("max", hasKey (F.max x x)) + , ("toMaybe", hasKey (F.toMaybe x)) + , ("fromMaybe", hasKey (F.fromMaybe 0 m)) + , ("isJust", hasKey (F.isJust m)) + , ("isNothing", hasKey (F.isNothing m)) + , ("isNull", hasKey (F.isNull m)) + , ("isNotNull", hasKey (F.isNotNull m)) + , ("fromJust", hasKey (F.fromJust m)) + , ("recode", hasKey (F.recode [(1 :: Int, 'a')] x)) + , ("recodeWithDefault", hasKey (F.recodeWithDefault 'z' [(1 :: Int, 'a')] x)) + , ("recodeWithCondition", hasKey (F.recodeWithCondition (F.lit (0 :: Int)) [((.> F.lit @Int 1), 1)] x)) + , ("firstOrNothing", hasKey (F.firstOrNothing l)) + , ("lastOrNothing", hasKey (F.lastOrNothing l)) + , ("contains", hasKey (F.contains "a" t)) + , ("length", hasKey (F.length t)) + , ("splitOn", hasKey (F.splitOn "," t)) + , ("match", hasKey (F.match "a+" t)) + , ("matchAll", hasKey (F.matchAll "a+" t)) + , ("parseDate", hasKey (F.parseDate @Day "%Y" t)) + , ("daysBetween", hasKey (F.daysBetween day day)) + , ("over", hasKey (F.over ["k"] (F.sum x))) + ] + +-- Functions taking arbitrary Haskell functions can only be anonymous. +higherOrderAreAnonymous :: Test +higherOrderAreAnonymous = + TestList + [ TestLabel fn (TestCase (assertBool fn (Prelude.not ok))) + | (fn, ok) <- anonymous + ] + where + x = F.col @Int "x" + m = F.col @(Maybe Int) "m" + anonymous = + [ ("lift", hasKey (F.lift (+ 1) x)) + , ("lift2", hasKey (F.lift2 (+) x x)) + , ("nullLift", hasKey (F.nullLift (+ 1) m)) + , ("nullLift2", hasKey (F.nullLift2 (+) m x)) + , ("reduce", hasKey (F.reduce x (0 :: Int) (+))) + , ("whenPresent", hasKey (F.whenPresent (+ 1) m)) + , ("whenBothPresent", hasKey (F.whenBothPresent (+) m m)) + , ("bind", hasKey (F.bind (Just . (+ 1)) m)) + ] + +tests :: [Test] +tests = + [ udfNamedLikeBuiltin + , udfOtherName + , anonymousNotEqual + , namedEqual + , namedVsBuiltin + , fromMaybeDefaults + , powVariants + , anonymousAtomsNotMerged + , fromMaybeTrueNotFalseOnNull + , momentsRespectUdf + , everyFunctionIsBuiltin + , higherOrderAreAnonymous + ] diff --git a/tests/Typed/Parity.hs b/tests/Typed/Parity.hs index 50c4d390..5d23a2de 100644 --- a/tests/Typed/Parity.hs +++ b/tests/Typed/Parity.hs @@ -12,6 +12,7 @@ import System.Random (mkStdGen) import qualified DataFrame as D import qualified DataFrame.Functions as F import qualified DataFrame.Internal.Column as DI +import DataFrame.Internal.Expression (eqExpr) import qualified DataFrame.Typed as DT import qualified DataFrame.Typed.Statistics as TS @@ -89,8 +90,29 @@ exprParity = "recodeWithDefault" (DT.recodeWithDefault 0 [(1 :: Int, 100 :: Int)] (DT.col @"x")) (F.recodeWithDefault 0 [(1 :: Int, 100 :: Int)] (F.col @Int "x")) + , deriveEq + "eq" + (DT.eq (DT.col @"x") (DT.lit 3)) + (F.eq (F.col @Int "x") (F.lit @Int 3)) + , deriveEq + "and" + (DT.and (DT.lt (DT.col @"x") (DT.lit 5)) (DT.geq (DT.col @"x") (DT.lit 2))) + (F.and (F.lt x (F.lit @Int 5)) (F.geq x (F.lit @Int 2))) + , deriveEq "contains" (DT.contains "a" (DT.col @"g")) (F.contains "a" g) + , deriveEq "length" (DT.length (DT.col @"g")) (F.length g) + , deriveEq + "isNull" + (DT.isNull (DT.toMaybe (DT.col @"x"))) + (F.isNull (F.toMaybe x)) + , deriveEq "cast" (DT.cast @"x" @Int) (F.cast @Int "x") + , deriveEq + "castWithDefault" + (DT.castWithDefault @"g" (0 :: Int)) + (F.castWithDefault (0 :: Int) "g") ] where + x = F.col @Int "x" + g = F.col @T.Text "g" deriveEq :: (DI.Columnable a) => String -> DT.TExpr S a -> D.Expr a -> Test @@ -102,6 +124,20 @@ exprParity = (DT.thaw (DT.derive @"r" te tdf)) ) +-- Typed aggregations build the same expression as their untyped forms. +aggParity :: Test +aggParity = + TestList + [ same "countDistinct" (DT.countDistinct (DT.col @"x")) (F.countDistinct x) + , same "topK" (DT.topK 2 (DT.col @"x")) (F.topK 2 x) + , same "bottomK" (DT.bottomK 2 (DT.col @"x")) (F.bottomK 2 x) + , same "not" (DT.not (DT.eq (DT.col @"x") 1)) (F.not (F.eq x 1)) + ] + where + x = F.col @Int "x" + same :: (DI.Columnable a) => String -> DT.TExpr S a -> D.Expr a -> Test + same lbl te ue = TestCase (assertBool lbl (eqExpr (DT.unTExpr te) ue)) + -- apply / valueCounts / horizontal merge. applyParity :: Test applyParity = @@ -199,6 +235,7 @@ tests :: [Test] tests = [ statParity , exprParity + , aggParity , applyParity , samplingParity , accessParity