Change window functions to operate on tables rather than columns

changelog.d/20231006_191956_shane.obrien_window_table.md

@@ -0,0 +1,3 @@
+### Changed
+
+- The window functions `lag`, `lead`, `firstValue`, `lastValue` and `nthValue` can now operate on entire rows at once as opposed to just single columns.

src/Rel8/Expr/Window.hs

@@ -6,11 +6,11 @@ module Rel8.Expr.Window
   , percentRank
   , cumeDist
   , ntile
-  , lag, lagOn
-  , lead, leadOn
-  , firstValue, firstValueOn
-  , lastValue, lastValueOn
-  , nthValue, nthValueOn
+  , lagExpr, lagExprOn
+  , leadExpr, leadExprOn
+  , firstValueExpr, firstValueExprOn
+  , lastValueExpr, lastValueExprOn
+  , nthValueExpr, nthValueExprOn
   )
 where
 
@@ -76,68 +76,68 @@ ntile buckets = fromWindowFunction $ fromPrimExpr . fromColumn <$>
 
 
 -- | [@lag(value, offset, default)@](https://www.postgresql.org/docs/current/functions-window.html)
-lag :: Expr Int32 -> Expr a -> Window (Expr a) (Expr a)
-lag offset def =
+lagExpr :: Expr Int32 -> Expr a -> Window (Expr a) (Expr a)
+lagExpr offset def =
   fromWindowFunction $
     dimap (toColumn . toPrimExpr) (fromPrimExpr . fromColumn) $
       Opaleye.lag (toColumn (toPrimExpr offset)) (toColumn (toPrimExpr def))
 
 
 -- | Applies 'lag' to the column selected by the given function.
-lagOn :: Expr Int32 -> Expr a -> (i -> Expr a) -> Window i (Expr a)
-lagOn offset def f = lmap f (lag offset def)
+lagExprOn :: Expr Int32 -> Expr a -> (i -> Expr a) -> Window i (Expr a)
+lagExprOn offset def f = lmap f (lagExpr offset def)
 
 
 -- | [@lead(value, offset, default)@](https://www.postgresql.org/docs/current/functions-window.html)
-lead :: Expr Int32 -> Expr a -> Window (Expr a) (Expr a)
-lead offset def =
+leadExpr :: Expr Int32 -> Expr a -> Window (Expr a) (Expr a)
+leadExpr offset def =
   fromWindowFunction $
     dimap (toColumn . toPrimExpr) (fromPrimExpr . fromColumn) $
       Opaleye.lead (toColumn (toPrimExpr offset)) (toColumn (toPrimExpr def))
 
 
 -- | Applies 'lead' to the column selected by the given function.
-leadOn :: Expr Int32 -> Expr a -> (i -> Expr a) -> Window i (Expr a)
-leadOn offset def f = lmap f (lead offset def)
+leadExprOn :: Expr Int32 -> Expr a -> (i -> Expr a) -> Window i (Expr a)
+leadExprOn offset def f = lmap f (leadExpr offset def)
 
 
 -- | [@first_value(value)@](https://www.postgresql.org/docs/current/functions-window.html)
-firstValue :: Window (Expr a) (Expr a)
-firstValue =
+firstValueExpr :: Window (Expr a) (Expr a)
+firstValueExpr =
   fromWindowFunction $
     dimap (toColumn . toPrimExpr) (fromPrimExpr . fromColumn)
       Opaleye.firstValue
 
 
 -- | Applies 'firstValue' to the column selected by the given function.
-firstValueOn :: (i -> Expr a) -> Window i (Expr a)
-firstValueOn f = lmap f firstValue
+firstValueExprOn :: (i -> Expr a) -> Window i (Expr a)
+firstValueExprOn f = lmap f firstValueExpr
 
 
 -- | [@last_value(value)@](https://www.postgresql.org/docs/current/functions-window.html)
-lastValue :: Window (Expr a) (Expr a)
-lastValue =
+lastValueExpr :: Window (Expr a) (Expr a)
+lastValueExpr =
   fromWindowFunction $
     dimap (toColumn . toPrimExpr) (fromPrimExpr . fromColumn)
       Opaleye.lastValue
 
 
 -- | Applies 'lastValue' to the column selected by the given function.
-lastValueOn :: (i -> Expr a) -> Window i (Expr a)
-lastValueOn f = lmap f lastValue
+lastValueExprOn :: (i -> Expr a) -> Window i (Expr a)
+lastValueExprOn f = lmap f lastValueExpr
 
 
 -- | [@nth_value(value, n)@](https://www.postgresql.org/docs/current/functions-window.html)
-nthValue :: Expr Int32 -> Window (Expr a) (Expr (Nullify a))
-nthValue n =
+nthValueExpr :: Expr Int32 -> Window (Expr a) (Expr (Nullify a))
+nthValueExpr n =
   fromWindowFunction $
     dimap (toColumn . toPrimExpr) (fromPrimExpr . fromColumn) $
       Opaleye.nthValue (toColumn (toPrimExpr n))
 
 
 -- | [@nth_value(value, n)@](https://www.postgresql.org/docs/current/functions-window.html)
-nthValueOn :: Expr Int32 -> (i -> Expr a) -> Window i (Expr (Nullify a))
-nthValueOn n f = lmap f (nthValue n)
+nthValueExprOn :: Expr Int32 -> (i -> Expr a) -> Window i (Expr (Nullify a))
+nthValueExprOn n f = lmap f (nthValueExpr n)
 
 
 fromAggregate :: Aggregator' fold i a -> Opaleye.Aggregator i a

src/Rel8/Schema/HTable.hs

@@ -16,7 +16,7 @@
 module Rel8.Schema.HTable
   ( HTable (HField, HConstrainTable)
   , hfield, htabulate, htraverse, hdicts, hspecs
-  , hfoldMap, hmap, htabulateA, htraverseP, htraversePWithField
+  , hfoldMap, hmap, htabulateA, htabulateP, htraverseP, htraversePWithField
   )
 where
 
@@ -136,6 +136,12 @@ htabulateA f = htraverse getCompose $ htabulate $ Compose . f
 {-# INLINABLE htabulateA #-}
 
 
+htabulateP :: (HTable t, ProductProfunctor p)
+  => (forall a. HField t a -> p i (context a)) -> p i (t context)
+htabulateP f = unApplyP $ htraverse (ApplyP . getCompose) $ htabulate $ Compose . f
+{-# INLINABLE htabulateP #-}
+
+
 newtype ApplyP p a b = ApplyP { unApplyP :: p a b }
 
 
@@ -154,8 +160,8 @@ htraverseP f = htraversePWithField (const f)
 
 htraversePWithField :: (HTable t, ProductProfunctor p)
   => (forall a. HField t a -> p (f a) (g a)) -> p (t f) (t g)
-htraversePWithField f = unApplyP $ htabulateA $ \field -> ApplyP $
-  lmap (flip hfield field) (f field)
+htraversePWithField f =
+  htabulateP $ \field -> lmap (flip hfield field) (f field)
 
 
 type GHField :: K.HTable -> Type -> Type

src/Rel8/Schema/HTable/Nullify.hs

@@ -32,9 +32,17 @@ import Data.Kind ( Type )
 import GHC.Generics ( Generic )
 import Prelude hiding ( null )
 
+-- profunctors
+import Data.Profunctor (lmap, rmap)
+
+-- product-profunctors
+import Data.Profunctor.Product (ProductProfunctor)
+
 -- rel8
 import Rel8.FCF ( Eval, Exp )
-import Rel8.Schema.HTable ( HTable, hfield, htabulate, htabulateA, hspecs )
+import Rel8.Schema.HTable
+  ( HTable, hfield, hspecs, htabulate, htabulateA, htabulateP
+  )
 import Rel8.Schema.HTable.MapTable
   ( HMapTable, HMapTableField( HMapTableField )
   , MapSpec, mapInfo
@@ -90,13 +98,11 @@ hnulls null = HNullify $ htabulate $ \(HMapTableField field) ->
 {-# INLINABLE hnulls #-}
 
 
-hnullify :: HTable t
-  => (forall a. Spec a -> context a -> context (Type.Nullify a))
-  -> t context
-  -> HNullify t context
-hnullify nullifier a = HNullify $ htabulate $ \(HMapTableField field) ->
-  case hfield hspecs field of
-    spec@Spec {} -> nullifier spec (hfield a field)
+hnullify :: (HTable t, ProductProfunctor p)
+  => (forall a. Spec a -> p (context a) (context (Type.Nullify a)))
+  -> p (t context) (HNullify t context)
+hnullify nullifier = rmap HNullify $ htabulateP $ \(HMapTableField field) ->
+  lmap (`hfield` field) (nullifier (hfield hspecs field))
 {-# INLINABLE hnullify #-}
 
 

src/Rel8/Table/Window.hs

@@ -1,20 +1,120 @@
+{-# language ApplicativeDo #-}
+{-# language FlexibleContexts #-}
 {-# language MonoLocalBinds #-}
+{-# language NamedFieldPuns #-}
 
 module Rel8.Table.Window
   ( currentRow
+  , lag, lagOn
+  , lead, leadOn
+  , firstValue, firstValueOn
+  , lastValue, lastValueOn
+  , nthValue, nthValueOn
   )
 where
 
 -- base
-import Prelude
+import Data.Int (Int32)
+import Prelude hiding (null)
 
 -- opaleye
 import qualified Opaleye.Window as Opaleye
 
+-- profunctor
+import Data.Profunctor (dimap, lmap)
+
 -- rel8
+import Rel8.Expr (Expr)
+import Rel8.Expr.Null (null, nullify, snull)
+import Rel8.Expr.Serialize (litExpr)
+import Rel8.Expr.Window
+  ( lagExpr, lagExprOn
+  , leadExpr, leadExprOn
+  , firstValueExpr
+  , lastValueExpr
+  , nthValueExpr, nthValueExprOn
+  )
+import Rel8.Schema.HTable (htraverseP)
+import Rel8.Schema.HTable.Identity (HIdentity (HIdentity))
+import Rel8.Schema.HTable.Label (hlabel)
+import Rel8.Schema.HTable.Maybe (HMaybeTable (HMaybeTable))
+import Rel8.Schema.HTable.Nullify (hnullify)
+import Rel8.Schema.Null (Nullity (NotNull, Null))
+import Rel8.Schema.Spec (Spec (..))
+import Rel8.Table (Table, fromColumns, toColumns)
+import Rel8.Table.Maybe (MaybeTable)
+import Rel8.Type.Tag (MaybeTag (IsJust))
 import Rel8.Window (Window (Window))
 
 
 -- | Return every column of the current row of a window query.
 currentRow :: Window a a
 currentRow = Window $ Opaleye.over (Opaleye.noWindowFunction id) mempty mempty
+
+
+-- | @'lag' n@ returns the row @n@ rows before the current row in a given
+-- window. Returns 'Rel8.nothingTable' if @n@ is out of bounds.
+lag :: Table Expr a => Expr Int32 -> Window a (MaybeTable Expr a)
+lag n = do
+  htag <- lagExprOn n null (\_ -> nullify (litExpr IsJust))
+  hjust <- lmap toColumns $ hnullify $ \Spec {info, nullity} ->
+    case nullity of
+      NotNull -> lagExprOn n (snull info) nullify
+      Null -> lagExpr n (snull info)
+  pure $ fromColumns $ HMaybeTable (hlabel (HIdentity htag)) (hlabel hjust)
+
+
+-- | Applies 'lag' to the columns selected by the given function.
+lagOn :: Table Expr a => Expr Int32 -> (i -> a) -> Window i (MaybeTable Expr a)
+lagOn n f = lmap f (lag n)
+
+
+-- | @'lead' n@ returns the row @n@ rows after the current row in a given
+-- window. Returns 'Rel8.nothingTable' if @n@ is out of bounds.
+lead :: Table Expr a => Expr Int32 -> Window a (MaybeTable Expr a)
+lead n = do
+  htag <- leadExprOn n null (\_ -> nullify (litExpr IsJust))
+  hjust <- lmap toColumns $ hnullify $ \Spec {info, nullity} ->
+    case nullity of
+      NotNull -> leadExprOn n (snull info) nullify
+      Null -> leadExpr n (snull info)
+  pure $ fromColumns $ HMaybeTable (hlabel (HIdentity htag)) (hlabel hjust)
+
+
+-- | Applies 'lead' to the columns selected by the given function.
+leadOn :: Table Expr a => Expr Int32 -> (i -> a) -> Window i (MaybeTable Expr a)
+leadOn n f = lmap f (lead n)
+
+
+-- | 'firstValue' returns the first row of the window of the current row.
+firstValue :: Table Expr a => Window a a
+firstValue = dimap toColumns fromColumns $ htraverseP firstValueExpr
+
+
+-- | Applies 'firstValue' to the columns selected by the given function.
+firstValueOn :: Table Expr a => (i -> a) -> Window i a
+firstValueOn f = lmap f firstValue
+
+
+-- | 'lastValue' returns the first row of the window of the current row.
+lastValue :: Table Expr a => Window a a
+lastValue = dimap toColumns fromColumns $ htraverseP lastValueExpr
+
+
+-- | Applies 'lastValue' to the columns selected by the given function.
+lastValueOn :: Table Expr a => (i -> a) -> Window i a
+lastValueOn f = lmap f lastValue
+
+
+-- | @'nthValue' n@ returns the @n@th row of the window of the current row.
+-- Returns 'Rel8.nothingTable' if @n@ is out of bounds.
+nthValue :: Table Expr a => Expr Int32 -> Window a (MaybeTable Expr a)
+nthValue n = do
+  htag <- nthValueExprOn n (\_ -> litExpr IsJust)
+  hjust <- lmap toColumns $ hnullify $ \_ -> nthValueExpr n
+  pure $ fromColumns $ HMaybeTable (hlabel (HIdentity htag)) (hlabel hjust)
+
+
+-- | Applies 'nthValue' to the columns selected by the given function.
+nthValueOn :: Table Expr a => Expr Int32 -> (i -> a) -> Window i (MaybeTable Expr a)
+nthValueOn n f = lmap f (nthValue n)