django/db/models/functions/window.py

from django.db.models.expressions import Func
from django.db.models.fields import FloatField, IntegerField

__all__ = [
    'CumeDist', 'DenseRank', 'FirstValue', 'Lag', 'LastValue', 'Lead',
    'NthValue', 'Ntile', 'PercentRank', 'Rank', 'RowNumber',
]


class CumeDist(Func):
    function = 'CUME_DIST'
    output_field = FloatField()
    window_compatible = True


class DenseRank(Func):
    function = 'DENSE_RANK'
    output_field = IntegerField()
    window_compatible = True


class FirstValue(Func):
    arity = 1
    function = 'FIRST_VALUE'
    window_compatible = True


class LagLeadFunction(Func):
    window_compatible = True

    def __init__(self, expression, offset=1, default=None, **extra):
        if expression is None:
            raise ValueError(
                '%s requires a non-null source expression.' %
                self.__class__.__name__
            )
        if offset is None or offset <= 0:
            raise ValueError(
                '%s requires a positive integer for the offset.' %
                self.__class__.__name__
            )
        args = (expression, offset)
        if default is not None:
            args += (default,)
        super().__init__(*args, **extra)

    def _resolve_output_field(self):
        sources = self.get_source_expressions()
        return sources[0].output_field


class Lag(LagLeadFunction):
    function = 'LAG'


class LastValue(Func):
    arity = 1
    function = 'LAST_VALUE'
    window_compatible = True


class Lead(LagLeadFunction):
    function = 'LEAD'


class NthValue(Func):
    function = 'NTH_VALUE'
    window_compatible = True

    def __init__(self, expression, nth=1, **extra):
        if expression is None:
            raise ValueError('%s requires a non-null source expression.' % self.__class__.__name__)
        if nth is None or nth <= 0:
            raise ValueError('%s requires a positive integer as for nth.' % self.__class__.__name__)
        super().__init__(expression, nth, **extra)

    def _resolve_output_field(self):
        sources = self.get_source_expressions()
        return sources[0].output_field


class Ntile(Func):
    function = 'NTILE'
    output_field = IntegerField()
    window_compatible = True

    def __init__(self, num_buckets=1, **extra):
        if num_buckets <= 0:
            raise ValueError('num_buckets must be greater than 0.')
        super().__init__(num_buckets, **extra)


class PercentRank(Func):
    function = 'PERCENT_RANK'
    output_field = FloatField()
    window_compatible = True


class Rank(Func):
    function = 'RANK'
    output_field = IntegerField()
    window_compatible = True


class RowNumber(Func):
    function = 'ROW_NUMBER'
    output_field = IntegerField()
    window_compatible = True