1from plain.models.expressions import Func
  2from plain.models.fields import FloatField, IntegerField
  3
  4__all__ = [
  5    "CumeDist",
  6    "DenseRank",
  7    "FirstValue",
  8    "Lag",
  9    "LastValue",
 10    "Lead",
 11    "NthValue",
 12    "Ntile",
 13    "PercentRank",
 14    "Rank",
 15    "RowNumber",
 16]
 17
 18
 19class CumeDist(Func):
 20    function = "CUME_DIST"
 21    output_field = FloatField()
 22    window_compatible = True
 23
 24
 25class DenseRank(Func):
 26    function = "DENSE_RANK"
 27    output_field = IntegerField()
 28    window_compatible = True
 29
 30
 31class FirstValue(Func):
 32    arity = 1
 33    function = "FIRST_VALUE"
 34    window_compatible = True
 35
 36
 37class LagLeadFunction(Func):
 38    window_compatible = True
 39
 40    def __init__(self, expression, offset=1, default=None, **extra):
 41        if expression is None:
 42            raise ValueError(
 43                f"{self.__class__.__name__} requires a non-null source expression."
 44            )
 45        if offset is None or offset <= 0:
 46            raise ValueError(
 47                f"{self.__class__.__name__} requires a positive integer for the offset."
 48            )
 49        args = (expression, offset)
 50        if default is not None:
 51            args += (default,)
 52        super().__init__(*args, **extra)
 53
 54    def _resolve_output_field(self):
 55        sources = self.get_source_expressions()
 56        return sources[0].output_field
 57
 58
 59class Lag(LagLeadFunction):
 60    function = "LAG"
 61
 62
 63class LastValue(Func):
 64    arity = 1
 65    function = "LAST_VALUE"
 66    window_compatible = True
 67
 68
 69class Lead(LagLeadFunction):
 70    function = "LEAD"
 71
 72
 73class NthValue(Func):
 74    function = "NTH_VALUE"
 75    window_compatible = True
 76
 77    def __init__(self, expression, nth=1, **extra):
 78        if expression is None:
 79            raise ValueError(
 80                f"{self.__class__.__name__} requires a non-null source expression."
 81            )
 82        if nth is None or nth <= 0:
 83            raise ValueError(
 84                f"{self.__class__.__name__} requires a positive integer as for nth."
 85            )
 86        super().__init__(expression, nth, **extra)
 87
 88    def _resolve_output_field(self):
 89        sources = self.get_source_expressions()
 90        return sources[0].output_field
 91
 92
 93class Ntile(Func):
 94    function = "NTILE"
 95    output_field = IntegerField()
 96    window_compatible = True
 97
 98    def __init__(self, num_buckets=1, **extra):
 99        if num_buckets <= 0:
100            raise ValueError("num_buckets must be greater than 0.")
101        super().__init__(num_buckets, **extra)
102
103
104class PercentRank(Func):
105    function = "PERCENT_RANK"
106    output_field = FloatField()
107    window_compatible = True
108
109
110class Rank(Func):
111    function = "RANK"
112    output_field = IntegerField()
113    window_compatible = True
114
115
116class RowNumber(Func):
117    function = "ROW_NUMBER"
118    output_field = IntegerField()
119    window_compatible = True