Avoiding/customizing COUNT(*) queries

[ethagnawl]

As mentioned elsewhere, a project I'm working on is very happily using this project and GDX. 😁

We are encountering an "issue" which is arguably tangential to what this project purports to do (optimized queryset construction and avoidance of N+1s) but does share overlapping concerns and is probably very attractive to its users.

The issue is the SELECT COUNT(*) which is automatically generated in different places by Graphene and (I believe) by this library. This query is potentially problematic because 1) its output may not be used and 2) it can be very costly on Postgres (~400ms in some of real-world scenarios).

This has been discussed at some length in the graphene-django issue tracker and there is a community/conventional "solution".

Unfortunately, I'm fairly certain the linked solution can't be used as a drop-in replacement in conjunction with this library. That solution may also break fundamental underpinnings of this library, like pagination. So, before I go down the path of dubious class extensions or forks of this project, I wanted to see if the maintainers or GDQO community had any thoughts about holistic approaches to this problem. Also, whether or not a PR would be considered.

Our use case is also slightly divergent from the linked one because, for performance reasons, we are also wanting to write custom SQL (which makes use of PG indexes) in the scenario where this query is required. So, I think an ideal solution would be some combination of the two (avoidance or optional customization).

As mentioned above, presumably this is a feature which others would find useful and it'd be great to have both that and my underlying assumption about whether this is even feasible in conjunction with this library validated.

Either way, I'll share what we learn here as we continue researching options.

Thanks!

[MrThearMan]

Let me try to summarize what you are asking:

1. You want to avoid running count(*) queries in connection fields when information related to pagination is not required (mainly totalCount and hasNextPage).
2. You'd also like a way to customize the step where count(*) happens without rewriting too much of the library.

For redusing count(*) queries, the linked solution seems to use a trick that works by fetching N+1 items from database to determine hasNextPage, although totalCount will still be wrong. Also, you couldn't use the last pagination argument either since calculate_queryset_slice requires size to be known so that last can be calculated. Given what we would be trying to do here, these seem like fine omissions, and seem doable to me.

There is already some conditional code for nested connection fields so that we don't do counts there if not necessary (those are much more expensive since they need to run subqueries), although this results in some of the worst code in the entire library. Modifying this to work with the above logic sounds awful, but I guess users would expect it to work similarly if top-level connections.

As for allowing customization, I think it would be a quite difficult given that pagination is split in two places (DjangoConnectionField.connection_resolver and QueryOptimizer.paginate_prefetch_queryset). The code here is also quite awful (as mentioned previously) so I think it would take considerable effort to wrangle it to any kind of stable interface.

[ethagnawl]

Yes, that is an accurate summary!

Fair enough about #2. I'm glad your read on the complexity of the undertaking squares with mine.

[rsomani95]

@MrThearMan thanks for your insight. I was able to somewhat wrangle together (1):

import query_optimizer
...


# Forward pagination only, following the implementation at
# https://gist.github.com/crucialfelix/d14506cdb9d3a04427a72aadcad30af4
def get_pagination_bounds(args: dict, max_limit: Optional[int] = None) -> dict:
    """
    Calculate pagination bounds without requiring total count.
    Only supports forward pagination (first/after).
    """

    if "before" in args:
        raise Exception("`before` not supported")
    if "last" in args:
        raise Exception("`last` not supported")

    pagination_args = validate_pagination_args(
        first=args.pop("first", None),
        last=None,  # Explicitly not supporting 'last'
        offset=args.pop("offset", None),
        after=args.pop("after", None),
        before=None,  # Explicitly not supporting 'before'
        max_limit=max_limit,
    )

    after = pagination_args.get("after")
    first = pagination_args.get("first", max_limit)

    start = after + 1 if after is not None else 0
    end = start + first + 1  # Get one extra to check for next page

    return dict(after=after, first=first, start=start, end=end)


class DjangoConnectionField(query_optimizer.DjangoConnectionField):
    def connection_resolver(
        self, root: Any, info: GQLInfo, **kwargs: Any
    ) -> ConnectionType:
        # Get pagination parameters, ignoring 'before' and 'last'
        pagination = get_pagination_bounds(kwargs.copy(), self.max_limit)

        # Store pagination info in context
        if not hasattr(info.context, "optimizer_pagination"):
            info.context.optimizer_pagination = {}
        name = to_snake_case(info.field_name)
        info.context.optimizer_pagination[name] = pagination

        # If field is aliased, a prefetch should have been done to that alias.
        # If not, call the ObjectType's "resolve_{field_name}" method, if it exists.
        # Otherwise, call the default resolver (usually `dict_or_attr_resolver`).
        alias = getattr(info.field_nodes[0].alias, "value", None)
        result = (
            getattr(root, alias)
            # Aliases don't matter at the root level, since we don't need to
            # distinguish them from a parent model prefetches.
            if root != info.root_value and alias is not None
            else self.resolver(root, info, **kwargs)
        )

        queryset = self.to_queryset(result)
        queryset = self.underlying_type.get_queryset(queryset, info)

        max_complexity = getattr(self.underlying_type._meta, "max_complexity", None)

        # Note if the queryset has already been optimized.
        already_optimized = is_optimized(queryset)

        optimizer = OptimizationCompiler(info, max_complexity=max_complexity).compile(
            queryset
        )
        # Optimise & filter the queryset first
        if optimizer is not None:
            queryset = optimizer.optimize_queryset(queryset)

        # Handle prefetched querysets
        if already_optimized:
            if isinstance(queryset, list):
                count = len(queryset)
            else:
                # If this is a nested connection field, prefetch queryset models should have been
                # annotated with the queryset count (pick it from the first one).
                count = getattr(
                    next(iter(getattr(queryset, "_result_cache", []) or []), None),
                    optimizer_settings.PREFETCH_COUNT_KEY,
                    0,  # QuerySet result cache is empty -> count is 0.
                )
            items = queryset
        # This is where we do the 'simplified' slicing
        else:
            # Slice queryset and fetch items
            sliced_queryset = queryset[pagination["start"] : pagination["end"]]
            items = list(fetch_in_context(sliced_queryset, info))

        # Check if there's a next page
        has_next_page = len(items) > pagination["first"]
        has_previous_page = pagination["start"] > 0

        # Trim to requested size
        items = items[: pagination["first"]]

        # Create edges
        edges: list[EdgeType] = [
            self.connection_type.Edge(
                node=value, cursor=offset_to_cursor(pagination["start"] + index)
            )
            for index, value in enumerate(items)
        ]

        connection = connection_adapter(
            cls=self.connection_type,
            edges=edges,
            pageInfo=page_info_adapter(
                startCursor=edges[0].cursor if edges else None,
                endCursor=edges[-1].cursor if edges else None,
                hasPreviousPage=has_previous_page,
                hasNextPage=has_next_page,
            ),
        )

        connection.iterable = queryset if already_optimized else sliced_queryset
        connection.length = count if already_optimized else len(items)
        return connection

I haven't tested this robustly, but as a first pass, it looks promising.

A few things come to mind:

• If totalCount is not being queried, and:
  • before or last are not part of the pagination args: this should be the default behavior as it's much faster (in my testing with postgres with a table of ~1.7 million rows)
  • If size can be provided by the user (hypothetically), before and last can also be supported