feat: Add fixity and precedence to binary operations.

Author: mchav

src/DataFrame/Functions.hs

@@ -44,6 +44,7 @@ import Text.Regex.TDFA
 import Prelude hiding (maximum, minimum)
 import Prelude as P
 
+infix 8 .^^, `div`
 infix 4 .==, .<, .<=, .>=, .>, ./=
 infixr 3 .&&
 infixr 2 .||
@@ -125,10 +126,10 @@ toDouble =
         )
 
 div :: (Integral a, Columnable a) => Expr a -> Expr a -> Expr a
-div = lift2Decorated Prelude.div "div" (Just "//") False 2
+div = lift2Decorated Prelude.div "div" (Just "//") False 7
 
 mod :: (Integral a, Columnable a) => Expr a -> Expr a -> Expr a
-mod = lift2Decorated Prelude.mod "mod" Nothing False 2
+mod = lift2Decorated Prelude.mod "mod" Nothing False 7
 
 (.==) :: (Columnable a, Eq a) => Expr a -> Expr a -> Expr Bool
 (.==) =
@@ -138,7 +139,7 @@ mod = lift2Decorated Prelude.mod "mod" Nothing False 2
             , binaryName = "eq"
             , binarySymbol = Just "=="
             , binaryCommutative = True
-            , binaryPrecedence = 1
+            , binaryPrecedence = 4
             }
         )
 
@@ -150,7 +151,7 @@ mod = lift2Decorated Prelude.mod "mod" Nothing False 2
             , binaryName = "neq"
             , binarySymbol = Just "/="
             , binaryCommutative = True
-            , binaryPrecedence = 1
+            , binaryPrecedence = 4
             }
         )
 
@@ -165,7 +166,7 @@ eq = (.==)
             , binaryName = "lt"
             , binarySymbol = Just "<"
             , binaryCommutative = False
-            , binaryPrecedence = 1
+            , binaryPrecedence = 4
             }
         )
 
@@ -180,7 +181,7 @@ lt = (.<)
             , binaryName = "gt"
             , binarySymbol = Just ">"
             , binaryCommutative = False
-            , binaryPrecedence = 1
+            , binaryPrecedence = 4
             }
         )
 
@@ -195,7 +196,7 @@ gt = (.>)
             , binaryName = "leq"
             , binarySymbol = Just "<="
             , binaryCommutative = False
-            , binaryPrecedence = 1
+            , binaryPrecedence = 4
             }
         )
 
@@ -210,7 +211,7 @@ leq = (.<=)
             , binaryName = "geq"
             , binarySymbol = Just ">="
             , binaryCommutative = False
-            , binaryPrecedence = 1
+            , binaryPrecedence = 4
             }
         )
 
@@ -228,7 +229,7 @@ and = (.&&)
             , binaryName = "and"
             , binarySymbol = Just "&&"
             , binaryCommutative = True
-            , binaryPrecedence = 1
+            , binaryPrecedence = 3
             }
         )
 
@@ -243,7 +244,7 @@ or = (.||)
             , binaryName = "or"
             , binarySymbol = Just "||"
             , binaryCommutative = True
-            , binaryPrecedence = 1
+            , binaryPrecedence = 2
             }
         )
 
@@ -328,10 +329,10 @@ zScore :: Expr Double -> Expr Double
 zScore c = (c - mean c) / stddev c
 
 pow :: (Columnable a, Num a) => Expr a -> Int -> Expr a
-pow _ 0 = Lit 1
-pow (Lit n) i = Lit (n ^ i)
-pow expr 1 = expr
-pow expr i = lift2Decorated (^) "pow" (Just "^") False 3 expr (lit i)
+pow expr i = lift2Decorated (^) "pow" (Just "^") False 8 expr (lit i)
+
+(.^^) :: (Columnable a, Num a) => Expr a -> Int -> Expr a
+(.^^) = pow
 
 relu :: (Columnable a, Num a, Ord a) => Expr a -> Expr a
 relu = lift (Prelude.max 0)

src/DataFrame/Internal/Expression.hs

@@ -68,7 +68,7 @@ instance (Num a, Columnable a) => Num (Expr a) where
                 , binaryName = "add"
                 , binarySymbol = Just "+"
                 , binaryCommutative = True
-                , binaryPrecedence = 1
+                , binaryPrecedence = 6
                 }
             )
 
@@ -80,7 +80,7 @@ instance (Num a, Columnable a) => Num (Expr a) where
                 , binaryName = "sub"
                 , binarySymbol = Just "-"
                 , binaryCommutative = False
-                , binaryPrecedence = 1
+                , binaryPrecedence = 6
                 }
             )
 
@@ -92,7 +92,7 @@ instance (Num a, Columnable a) => Num (Expr a) where
                 , binaryName = "mult"
                 , binarySymbol = Just "*"
                 , binaryCommutative = True
-                , binaryPrecedence = 2
+                , binaryPrecedence = 7
                 }
             )
 
@@ -133,7 +133,7 @@ instance (Fractional a, Columnable a) => Fractional (Expr a) where
                 , binaryName = "divide"
                 , binarySymbol = Just "/"
                 , binaryCommutative = True
-                , binaryPrecedence = 2
+                , binaryPrecedence = 7
                 }
             )
 
@@ -160,7 +160,7 @@ instance (Floating a, Columnable a) => Floating (Expr a) where
                 , binaryName = "exponentiate"
                 , binarySymbol = Just "**"
                 , binaryCommutative = False
-                , binaryPrecedence = 3
+                , binaryPrecedence = 8
                 }
             )
     log :: (Floating a, Columnable a) => Expr a -> Expr a
@@ -338,22 +338,37 @@ getColumns (Binary op l r) = getColumns l <> getColumns r
 getColumns (Agg strategy expr) = getColumns expr
 
 prettyPrint :: Expr a -> String
-prettyPrint = go 0
+prettyPrint = go 0 0
   where
-    go :: Int -> Expr a -> String
-    go prec expr = case expr of
+    indent :: Int -> String
+    indent n = replicate (n * 2) ' '
+
+    go :: Int -> Int -> Expr a -> String
+    go depth prec expr = case expr of
         Col name -> T.unpack name
         Lit value -> show value
         If cond t e ->
-            "if (" ++ go 0 cond ++ ") then (" ++ go 0 t ++ ") else (" ++ go 0 e ++ ")"
+            let inner =
+                    "if "
+                        ++ go (depth + 1) 0 cond
+                        ++ "\n"
+                        ++ indent (depth + 1)
+                        ++ "then "
+                        ++ go (depth + 1) 0 t
+                        ++ "\n"
+                        ++ indent (depth + 1)
+                        ++ "else "
+                        ++ go (depth + 1) 0 e
+             in if prec > 0 then "(" ++ inner ++ ")" else inner
         Unary op arg -> case unarySymbol op of
-            Nothing -> T.unpack (unaryName op) ++ "(" ++ go 0 arg ++ ")"
-            Just sym -> T.unpack sym ++ "(" ++ go 0 arg ++ ")"
+            Nothing -> T.unpack (unaryName op) ++ "(" ++ go depth 0 arg ++ ")"
+            Just sym -> T.unpack sym ++ "(" ++ go depth 0 arg ++ ")"
         Binary op l r ->
             let p = binaryPrecedence op
                 inner = case binarySymbol op of
-                    Just name -> go p l ++ " " ++ T.unpack name ++ " " ++ go p r
-                    Nothing -> T.unpack (binaryName op) ++ "(" ++ go p l ++ ", " ++ go p r ++ ")"
+                    Just name -> go depth p l ++ " " ++ T.unpack name ++ " " ++ go depth p r
+                    Nothing ->
+                        T.unpack (binaryName op) ++ "(" ++ go depth p l ++ ", " ++ go depth p r ++ ")"
              in if prec > p then "(" ++ inner ++ ")" else inner
-        Agg (CollectAgg op _) arg -> T.unpack op ++ "(" ++ go 0 arg ++ ")"
-        Agg (FoldAgg op _ _) arg -> T.unpack op ++ "(" ++ go 0 arg ++ ")"
+        Agg (CollectAgg op _) arg -> T.unpack op ++ "(" ++ go depth 0 arg ++ ")"
+        Agg (FoldAgg op _ _) arg -> T.unpack op ++ "(" ++ go depth 0 arg ++ ")"