Window functions

rel8.cabal

@@ -82,6 +82,7 @@ library
     Rel8.Expr.Order
     Rel8.Expr.Sequence
     Rel8.Expr.Serialize
+    Rel8.Expr.Window
 
     Rel8.FCF
 
@@ -120,6 +121,7 @@ library
     Rel8.Query.SQL
     Rel8.Query.These
     Rel8.Query.Values
+    Rel8.Query.Window
 
     Rel8.Schema.Context.Nullify
     Rel8.Schema.Dict
@@ -179,6 +181,7 @@ library
     Rel8.Table.These
     Rel8.Table.Transpose
     Rel8.Table.Undefined
+    Rel8.Table.Window
 
     Rel8.Type
     Rel8.Type.Array
@@ -197,6 +200,9 @@ library
     Rel8.Type.Sum
     Rel8.Type.Tag
 
+    Rel8.Window
+
+
 test-suite tests
   type:             exitcode-stdio-1.0
   build-depends:

src/Rel8.hs

@@ -274,6 +274,24 @@ module Rel8
   , nullsLast
 
     -- ** Window functions
+  , Window
+  , window
+  , Partition
+  , over
+  , partitionBy
+  , orderPartitionBy
+  , cumulative
+  , rowNumber
+  , rank
+  , denseRank
+  , percentRank
+  , cumeDist
+  , ntile
+  , lag
+  , lead
+  , firstValue
+  , lastValue
+  , nthValue
   , indexed
 
     -- ** Bindings
@@ -346,6 +364,7 @@ import Rel8.Expr.Order
 import Rel8.Expr.Serialize
 import Rel8.Expr.Sequence
 import Rel8.Expr.Text ( like, ilike )
+import Rel8.Expr.Window hiding ( cumulative )
 import Rel8.Generic.Rel8able ( KRel8able, Rel8able )
 import Rel8.Order
 import Rel8.Query
@@ -367,6 +386,7 @@ import Rel8.Query.SQL (showQuery)
 import Rel8.Query.Set
 import Rel8.Query.These
 import Rel8.Query.Values
+import Rel8.Query.Window
 import Rel8.Schema.Field
 import Rel8.Schema.HTable
 import Rel8.Schema.Name
@@ -402,6 +422,7 @@ import Rel8.Table.Rel8able ()
 import Rel8.Table.Serialize
 import Rel8.Table.These
 import Rel8.Table.Transpose
+import Rel8.Table.Window
 import Rel8.Type
 import Rel8.Type.Composite
 import Rel8.Type.Eq
@@ -416,6 +437,7 @@ import Rel8.Type.ReadShow
 import Rel8.Type.Semigroup
 import Rel8.Type.String
 import Rel8.Type.Sum
+import Rel8.Window
 
 
 -- $running

src/Rel8/Expr/Window.hs

@@ -0,0 +1,96 @@
+module Rel8.Expr.Window
+  ( cumulative
+  , rowNumber
+  , rank
+  , denseRank
+  , percentRank
+  , cumeDist
+  , ntile
+  , lag
+  , lead
+  , firstValue
+  , lastValue
+  , nthValue
+  )
+where
+
+-- base
+import Data.Int ( Int32, Int64 )
+import Prelude
+
+-- opaleye
+import qualified Opaleye.Window as Opaleye
+
+-- rel8
+import Rel8.Aggregate ( Aggregate( Aggregate ) )
+import Rel8.Expr ( Expr )
+import Rel8.Expr.Opaleye ( fromColumn, fromPrimExpr, toColumn, toPrimExpr )
+import Rel8.Schema.Null ( Nullify )
+import Rel8.Window ( Window( Window ) )
+
+
+cumulative :: Aggregate a -> Window (Expr a)
+cumulative (Aggregate aggregator) = Window $ Opaleye.cumulative aggregator ()
+
+
+-- | [@row_number()@](https://www.postgresql.org/docs/current/functions-window.html)
+rowNumber :: Window (Expr Int64)
+rowNumber = Window $ fromPrimExpr . fromColumn <$> Opaleye.rowNumber
+
+
+-- | [@rank()@](https://www.postgresql.org/docs/current/functions-window.html)
+rank :: Window (Expr Int64)
+rank = Window $ fromPrimExpr . fromColumn <$> Opaleye.rank
+
+
+-- | [@dense_rank()@](https://www.postgresql.org/docs/current/functions-window.html)
+denseRank :: Window (Expr Int64)
+denseRank = Window $ fromPrimExpr . fromColumn <$> Opaleye.denseRank
+
+
+-- | [@percent_rank()@](https://www.postgresql.org/docs/current/functions-window.html)
+percentRank :: Window (Expr Double)
+percentRank = Window $ fromPrimExpr . fromColumn <$> Opaleye.percentRank
+
+
+-- | [@cume_dist()@](https://www.postgresql.org/docs/current/functions-window.html)
+cumeDist :: Window (Expr Double)
+cumeDist = Window $ fromPrimExpr . fromColumn <$> Opaleye.cumeDist
+
+
+-- | [@ntile(num_buckets)@](https://www.postgresql.org/docs/current/functions-window.html)
+ntile :: Expr Int32 -> Window (Expr Int32)
+ntile buckets = Window $ fromPrimExpr . fromColumn <$>
+  Opaleye.ntile (toColumn (toPrimExpr buckets))
+
+
+-- | [@lag(value, offset, default)@](https://www.postgresql.org/docs/current/functions-window.html)
+lag :: Expr a -> Expr Int32 -> Expr a -> Window (Expr a)
+lag a offset def = Window $ fromPrimExpr . fromColumn <$>
+  Opaleye.lag (toColumn (toPrimExpr a)) (toColumn (toPrimExpr offset))
+    (toColumn (toPrimExpr def))
+
+
+-- | [@lead(value, offset, default)@](https://www.postgresql.org/docs/current/functions-window.html)
+lead :: Expr a -> Expr Int32 -> Expr a -> Window (Expr a)
+lead a offset def = Window $ fromPrimExpr . fromColumn <$>
+  Opaleye.lead (toColumn (toPrimExpr a)) (toColumn (toPrimExpr offset))
+    (toColumn (toPrimExpr def))
+
+
+-- | [@first_value(value)@](https://www.postgresql.org/docs/current/functions-window.html)
+firstValue :: Expr a -> Window (Expr a)
+firstValue a = Window $ fromPrimExpr . fromColumn <$>
+  Opaleye.firstValue (toColumn (toPrimExpr a))
+
+
+-- | [@last_value(value)@](https://www.postgresql.org/docs/current/functions-window.html)
+lastValue :: Expr a -> Window (Expr a)
+lastValue a = Window $ fromPrimExpr . fromColumn <$>
+  Opaleye.lastValue (toColumn (toPrimExpr a))
+
+
+-- | [@nth_value(value, n)@](https://www.postgresql.org/docs/current/functions-window.html)
+nthValue :: Expr a -> Expr Int32 -> Window (Expr (Nullify a))
+nthValue a n = Window $ fromPrimExpr . fromColumn <$>
+  Opaleye.nthValue (toColumn (toPrimExpr a)) (toColumn (toPrimExpr n))

src/Rel8/Query/Indexed.hs

@@ -4,31 +4,17 @@ module Rel8.Query.Indexed
 where
 
 -- base
+import Control.Applicative ( liftA2 )
 import Data.Int ( Int64 )
 import Prelude
 
--- opaleye
-import qualified Opaleye.Internal.HaskellDB.PrimQuery as Opaleye
-import qualified Opaleye.Internal.PackMap as Opaleye
-import qualified Opaleye.Internal.PrimQuery as Opaleye
-import qualified Opaleye.Internal.QueryArr as Opaleye
-import qualified Opaleye.Internal.Tag as Opaleye
-
 -- rel8
 import Rel8.Expr ( Expr )
-import Rel8.Expr.Opaleye ( fromPrimExpr )
+import Rel8.Expr.Window ( rowNumber )
 import Rel8.Query ( Query )
-import Rel8.Query.Opaleye ( mapOpaleye )
+import Rel8.Query.Window ( window )
 
 
 -- | Pair each row of a query with its index within the query.
 indexed :: Query a -> Query (Expr Int64, a)
-indexed = mapOpaleye $ \f -> Opaleye.stateQueryArr $ \_ tag ->
-  let
-    (a, query, tag') = Opaleye.runStateQueryArr f () tag
-    tag'' = Opaleye.next tag'
-    window = Opaleye.ConstExpr $ Opaleye.OtherLit "ROW_NUMBER() OVER () - 1"
-    (index, bindings) = Opaleye.run $ Opaleye.extractAttr "index" tag' window
-    query' = query <> Opaleye.aRebind bindings
-  in
-    ((fromPrimExpr index, a), query', tag'')
+indexed = window . fmap (liftA2 (,) rowNumber . pure)

src/Rel8/Query/Window.hs

@@ -0,0 +1,26 @@
+module Rel8.Query.Window
+  ( window
+  )
+where
+
+-- base
+import Prelude
+
+-- opaleye
+import qualified Opaleye.Window as Opaleye
+
+-- rel8
+import Rel8.Query ( Query )
+import Rel8.Query.Opaleye ( mapOpaleye )
+import Rel8.Window ( Window( Window ) )
+
+
+-- | 'window' runs a query composed of expressions containing
+-- [window functions](https://www.postgresql.org/docs/current/tutorial-window.html).
+-- 'window' is similar to 'Rel8.aggregate', with the main difference being
+-- that in a window query, each input row corresponds to one output row,
+-- whereas aggregation queries fold the entire input query down into a single
+-- row. To put this into a Haskell context, 'Rel8.aggregate' is to 'foldl' as
+-- 'window' is to 'scanl'.
+window :: Query (Window a) -> Query a
+window = mapOpaleye (Opaleye.window . fmap (\(Window a) -> a))

src/Rel8/Schema/HTable.hs

@@ -16,11 +16,12 @@
 module Rel8.Schema.HTable
   ( HTable (HField, HConstrainTable)
   , hfield, htabulate, htraverse, hdicts, hspecs
-  , hmap, htabulateA
+  , hfoldMap, hmap, htabulateA
   )
 where
 
 -- base
+import Data.Functor.Const ( Const( Const ), getConst )
 import Data.Kind ( Constraint, Type )
 import Data.Functor.Compose ( Compose( Compose ), getCompose )
 import Data.Proxy ( Proxy )
@@ -114,6 +115,11 @@ class HTable t where
   {-# INLINABLE hspecs #-}
 
 
+hfoldMap :: (HTable t, Semigroup s)
+  => (forall a. context a -> s) -> t context -> s
+hfoldMap f a = getConst $ htraverse (Const . f) a
+
+
 hmap :: HTable t
   => (forall a. context a -> context' a) -> t context -> t context'
 hmap f a = htabulate $ \field -> f (hfield a field)

src/Rel8/Table/Window.hs

@@ -0,0 +1,24 @@
+{-# language MonoLocalBinds #-}
+
+module Rel8.Table.Window
+  ( cumulative
+  )
+where
+
+-- base
+import Prelude
+
+-- rel8
+import Rel8.Aggregate ( Aggregates )
+import qualified Rel8.Expr.Window as Expr
+import Rel8.Schema.HTable ( htraverse )
+import Rel8.Table ( fromColumns, toColumns )
+import Rel8.Window ( Window )
+
+
+-- | 'cumulative' allows the use of aggregation functions in 'Window'
+-- expressions. In particular, @'cumulative' . 'Rel8.sum'@
+-- (when combined with 'Rel8.Window.orderPartitionBy') gives a running total,
+-- also known as a \"cumulative sum\", hence the name @cumulative@.
+cumulative :: Aggregates aggregates exprs => aggregates -> Window exprs
+cumulative = fmap fromColumns . htraverse Expr.cumulative . toColumns

src/Rel8/Window.hs

@@ -0,0 +1,78 @@
+{-# language DerivingVia #-}
+{-# language FlexibleContexts #-}
+{-# language GeneralizedNewtypeDeriving #-}
+{-# language StandaloneKindSignatures #-}
+
+module Rel8.Window
+  ( Window(..)
+  , Partition
+  , over
+  , partitionBy
+  , orderPartitionBy
+  )
+where
+
+-- base
+import Data.Kind ( Type )
+import Prelude
+
+-- opaleye
+import qualified Opaleye.Internal.Window as Opaleye
+import qualified Opaleye.Internal.PackMap as Opaleye
+
+-- rel8
+import Rel8.Expr ( Expr )
+import Rel8.Expr.Opaleye ( toColumn, toPrimExpr )
+import Rel8.Order( Order( Order ) )
+import Rel8.Schema.HTable ( hfoldMap )
+import Rel8.Table ( Table, toColumns )
+
+-- semigroupoids
+import Data.Functor.Apply ( Apply, WrappedApplicative(..) )
+
+
+-- | 'Window' is an applicative functor that represents expressions that
+-- contain
+-- [window functions](https://www.postgresql.org/docs/current/tutorial-window.html).
+-- 'Rel8.Query.Window.window' can be used to
+-- evaluate these expressions over a particular query.
+type Window :: Type -> Type
+newtype Window a = Window (Opaleye.Window a)
+  deriving newtype (Functor, Applicative)
+  deriving (Apply) via (WrappedApplicative Window)
+
+
+-- | In PostgreSQL, window functions must specify the \"window\" or
+-- \"partition\" over which they operate. The syntax for this looks like:
+-- @SUM(salary) OVER (PARTITION BY department)@. The Rel8 type 'Partition'
+-- represents everything that comes after @OVER@.
+--
+-- 'Partition' is a 'Monoid', so 'Partition's created with 'partitionBy' and
+-- 'orderPartitionBy' can be combined using '<>'.
+type Partition :: Type
+newtype Partition = Partition Opaleye.Partition
+  deriving newtype (Semigroup, Monoid)
+
+
+-- | 'over' adds a 'Partition' to a 'Window' expression.
+--
+-- @@@
+-- 'Rel8.Table.Window.cumulative' ('Rel8.Expr.Aggregate.sum' salary) `over` 'partitionBy' department <> 'orderPartitionBy' salary 'Rel8.desc'
+-- @@@
+over :: Window a -> Partition -> Window a
+over (Window w) (Partition p) = Window (w `Opaleye.over` p)
+infixl 1 `over`
+
+
+-- | Restricts a window function to operate only the group of rows that share
+-- the same value(s) for the given expression(s).
+partitionBy :: Table Expr a => a -> Partition
+partitionBy = Partition . hfoldMap opartitionBy . toColumns
+  where
+    opartitionBy = Opaleye.partitionBy . toColumn . toPrimExpr
+
+
+-- | Controls the order in which rows are processed by window functions. This
+-- does not need to match the ordering of the overall query.
+orderPartitionBy :: a -> Order a -> Partition
+orderPartitionBy a (Order ordering) = Partition $ Opaleye.orderPartitionBy a ordering