<commit_before>from __future__ import absolute_import
from django.contrib.postgres.fields.array import ArrayField
from django.db.models import CharField, IntegerField
from django.db.models.aggregates import Max
from django.db.models.expressions import F, Func, Value
from django.db.models.query import Q, QuerySet
from django_cte import With
from mptt.models import MPTTModel, TreeManager

from .queryutil import ComparedQuerySet, TimingContext

int_field = IntegerField()
str_array = ArrayField(CharField())


class StrArray(Func):
    function = "Array"
    # HACK fool postgres with concat
    # https://stackoverflow.com/a/12488455/10840 (see comment by KajMagnus)
    template = "%(function)s[%(expressions)s || '']::varchar[]"
    output_field = str_array


class array_append(Func):
    function = "array_append"
    output_field = str_array


class array_length(Func):
    function = "array_length"
    template = "%(function)s(%(expressions)s, 1)"
    output_field = int_field


class AdjListManager(TreeManager):

    def cte_get_ancestors(self, node, ascending=False, include_self=False):
        """Query node ancestors

        :param node: A model instance or a QuerySet or Q object querying
        the adjacency list model. If a QuerySet, it should query a
        single value with something like `.values('id')`. If Q the
        `include_self` argument will be ignored.
        :param ascending: Order of results. The default (`False`) gets
        results in descending order (root ancestor first, immediate
        parent last).
        :param include_self:
        :returns: A `QuerySet` instance.
        """
        parent_col = self.model.parent_id_attr

        if isinstance(node, Q):
            where = node
        elif include_self:
            if isinstance(node, QuerySet):
                where = Q(id__in=node)
            else:
                where = Q(id=node.id)
        elif isinstance(node, QuerySet):
            where = Q(id__in=node.values(parent_col))
        else:
            where = Q(id=getattr(node, parent_col))

        def make_cte_query(cte):
            return self.filter(where).values(
                "id",
                parent_col,
                _depth=Value(0, output_field=int_field),
            ).union(
                cte.join(self.model, id=getattr(cte.col, parent_col)).values(
                    "id",
                    parent_col,
                    _depth=cte.col._depth + Value(1, output_field=int_field),
                ),
            )

        cte = With.recursive(make_cte_query)
        return (
            cte
            .join(self.all(), id=cte.col.id)
            .with_cte(cte)
            .order_by(("" if ascending else "-") + "{}._depth".format(cte.name))
        )

    def cte_get_descendants(self, node, include_self=False):
        """Query node descendants

        :param node: A model instance or a QuerySet or Q object querying
        the adjacency list model. If a QuerySet, it should query a
        single value with something like `.values('id')`. If Q the
        `include_self` argument will be ignored.
        :returns: A `QuerySet` instance.
        """
        parent_col = self.model.parent_id_attr
        ordering_col = self.model.ordering_col_attr

        discard_dups = False
        if isinstance(node, Q):
            where = node
            discard_dups = True
        elif include_self:
            if isinstance(node, QuerySet):
                where = Q(id__in=node)
                discard_dups = True
            else:
                where = Q(id=node.id)
        elif isinstance(node, QuerySet):
            where = Q(**{parent_col + "__in": node})
            discard_dups = True
        else:
            where = Q(**{parent_col: node.id})

        def make_cte_query(cte):
            return self.filter(where).values(
                "id",
                parent_col,
                _cte_ordering=StrArray(ordering_col),
            ).union(
                cte.join(self.model, **{parent_col: cte.col.id}).annotate(
                    _cte_ordering=array_append(
                        cte.col._cte_ordering,
                        F(ordering_col),
                    )
                ).values(
                    "id",
                    parent_col,
                    "_cte_ordering",
                ),
                all=True,
            )
        cte = With.recursive(make_cte_query)<commit_msg>Remove ORDER BY from internal CTE queries<commit_after>from __future__ import absolute_import
from django.contrib.postgres.fields.array import ArrayField
from django.db.models import CharField, IntegerField
from django.db.models.aggregates import Max
from django.db.models.expressions import F, Func, Value
from django.db.models.query import Q, QuerySet
from django_cte import With
from mptt.models import MPTTModel, TreeManager

from .queryutil import ComparedQuerySet, TimingContext

int_field = IntegerField()
str_array = ArrayField(CharField())


class StrArray(Func):
    function = "Array"
    # HACK fool postgres with concat
    # https://stackoverflow.com/a/12488455/10840 (see comment by KajMagnus)
    template = "%(function)s[%(expressions)s || '']::varchar[]"
    output_field = str_array


class array_append(Func):
    function = "array_append"
    output_field = str_array


class array_length(Func):
    function = "array_length"
    template = "%(function)s(%(expressions)s, 1)"
    output_field = int_field


class AdjListManager(TreeManager):

    def cte_get_ancestors(self, node, ascending=False, include_self=False):
        """Query node ancestors

        :param node: A model instance or a QuerySet or Q object querying
        the adjacency list model. If a QuerySet, it should query a
        single value with something like `.values('id')`. If Q the
        `include_self` argument will be ignored.
        :param ascending: Order of results. The default (`False`) gets
        results in descending order (root ancestor first, immediate
        parent last).
        :param include_self:
        :returns: A `QuerySet` instance.
        """
        parent_col = self.model.parent_id_attr

        if isinstance(node, Q):
            where = node
        elif include_self:
            if isinstance(node, QuerySet):
                where = Q(id__in=node.order_by())
            else:
                where = Q(id=node.id)
        elif isinstance(node, QuerySet):
            where = Q(id__in=node.order_by().values(parent_col))
        else:
            where = Q(id=getattr(node, parent_col))

        def make_cte_query(cte):
            return self.filter(where).order_by().values(
                "id",
                parent_col,
                _depth=Value(0, output_field=int_field),
            ).union(
                cte.join(
                    self.all().order_by(),
                    id=getattr(cte.col, parent_col)
                ).values(
                    "id",
                    parent_col,
                    _depth=cte.col._depth + Value(1, output_field=int_field),
                ),
            )

        cte = With.recursive(make_cte_query)
        return (
            cte
            .join(self.all(), id=cte.col.id)
            .with_cte(cte)
            .order_by(("" if ascending else "-") + "{}._depth".format(cte.name))
        )

    def cte_get_descendants(self, node, include_self=False):
        """Query node descendants

        :param node: A model instance or a QuerySet or Q object querying
        the adjacency list model. If a QuerySet, it should query a
        single value with something like `.values('id')`. If Q the
        `include_self` argument will be ignored.
        :returns: A `QuerySet` instance.
        """
        parent_col = self.model.parent_id_attr
        ordering_col = self.model.ordering_col_attr

        discard_dups = False
        if isinstance(node, Q):
            where = node
            discard_dups = True
        elif include_self:
            if isinstance(node, QuerySet):
                where = Q(id__in=node.order_by())
                discard_dups = True
            else:
                where = Q(id=node.id)
        elif isinstance(node, QuerySet):
            where = Q(**{parent_col + "__in": node.order_by()})
            discard_dups = True
        else:
            where = Q(**{parent_col: node.id})

        def make_cte_query(cte):
            return self.filter(where).order_by().values(
                "id",
                parent_col,
                _cte_ordering=StrArray(ordering_col),
            ).union(
                cte.join(
                    self.all().order_by(),
                    **{parent_col: cte.col.id}
                ).annotate(
                    _cte_ordering=array_append(
                        cte.col._cte_ordering,
                        F(ordering_col),
                    )
                ).values(
                    "id",
                    parent_col,
                    "_cte_ordering",
                ),
                all=True,
            )
        cte = With.recursive(make_cte_query)
{'fraction_non_alphanumeric': 0.09622972380534853, 'fraction_numerical': 0.0035072336694432268, 'mean_word_length': 1.8524538918412004, 'pattern_counts': {'":': 2, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1067238', 'n_tokens_mistral': 2694, 'n_tokens_neox': 2539, 'n_words': 637}