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                "%s requires a non-null source expression." % self.__class__.__name__
 44            )
 45        if offset is None or offset <= 0:
 46            raise ValueError(
 47                "%s requires a positive integer for the offset."
 48                % self.__class__.__name__
 49            )
 50        args = (expression, offset)
 51        if default is not None:
 52            args += (default,)
 53        super().__init__(*args, **extra)
 54
 55    def _resolve_output_field(self):
 56        sources = self.get_source_expressions()
 57        return sources[0].output_field
 58
 59
 60class Lag(LagLeadFunction):
 61    function = "LAG"
 62
 63
 64class LastValue(Func):
 65    arity = 1
 66    function = "LAST_VALUE"
 67    window_compatible = True
 68
 69
 70class Lead(LagLeadFunction):
 71    function = "LEAD"
 72
 73
 74class NthValue(Func):
 75    function = "NTH_VALUE"
 76    window_compatible = True
 77
 78    def __init__(self, expression, nth=1, **extra):
 79        if expression is None:
 80            raise ValueError(
 81                "%s requires a non-null source expression." % self.__class__.__name__
 82            )
 83        if nth is None or nth <= 0:
 84            raise ValueError(
 85                "%s requires a positive integer as for nth." % self.__class__.__name__
 86            )
 87        super().__init__(expression, nth, **extra)
 88
 89    def _resolve_output_field(self):
 90        sources = self.get_source_expressions()
 91        return sources[0].output_field
 92
 93
 94class Ntile(Func):
 95    function = "NTILE"
 96    output_field = IntegerField()
 97    window_compatible = True
 98
 99    def __init__(self, num_buckets=1, **extra):
100        if num_buckets <= 0:
101            raise ValueError("num_buckets must be greater than 0.")
102        super().__init__(num_buckets, **extra)
103
104
105class PercentRank(Func):
106    function = "PERCENT_RANK"
107    output_field = FloatField()
108    window_compatible = True
109
110
111class Rank(Func):
112    function = "RANK"
113    output_field = IntegerField()
114    window_compatible = True
115
116
117class RowNumber(Func):
118    function = "ROW_NUMBER"
119    output_field = IntegerField()
120    window_compatible = True