PyMongo Aggregation Pipeline: Practical Usage and Code

When you need to transform, filter, and aggregate MongoDB documents from a Python application, the aggregation pipeline is the tool to use. This article shows how to build and run MongoDB aggregation pipelines with PyMongo, covering essential stages, performance considerations, and common pitfalls. You'll learn how to construct pipelines that run efficiently and return the exact shape of data your application needs.

Understanding the Aggregation Pipeline

The aggregation pipeline in MongoDB processes documents through a sequence of stages. Each stage transforms the input documents and passes the result to the next stage. This is similar to Unix pipes: the output of one command becomes the input to the next. In PyMongo, you define the pipeline as a list of dictionaries, each representing a stage. The aggregate() method on a collection executes the pipeline and returns a cursor.

A simple pipeline might start with a $match stage to filter documents, followed by a $group stage to compute aggregates, and then a $sort stage to order the results. The order of stages matters because each stage operates on the documents produced by the previous one. For example, filtering early with $match reduces the number of documents that subsequent stages need to process, which can significantly improve performance.

Setting Up a Basic Aggregation in PyMongo

To run an aggregation, you need a collection object and a pipeline list. Here's a minimal example that counts documents by a field:

from pymongo import MongoClient

client = MongoClient('mongodb://localhost:27017')
db = client['sales_db']
collection = db['orders']

pipeline = [
    {'$group': {'_id': '$status', 'count': {'$sum': 1}}}
]

result = collection.aggregate(pipeline)
for doc in result:
    print(doc)

This pipeline groups all orders by the status field and counts how many documents share each status. The $group stage requires an _id field that defines the grouping key. The count field uses $sum to accumulate the number of documents in each group. The result is a cursor that yields documents like {'_id': 'shipped', 'count': 42}.

You can also pass options to aggregate(), such as batchSize to control how many documents are fetched from the server at once. This is useful when you expect a large result set and want to avoid loading everything into memory.

Common Aggregation Stages in PyMongo

PyMongo doesn't add its own stage syntax; you use the same aggregation stage operators as the MongoDB server. Here are the stages you'll use most often:

$match: Filters documents using a query expression. Place it early to reduce the data volume.
$project: Reshapes documents by including, excluding, or adding fields. You can compute new fields using expressions.
$group: Groups documents by a key and applies accumulator expressions like $sum, $avg, $min, $max, and $push.
$sort: Orders documents by one or more fields. This can be memory-intensive if applied to a large result set without a supporting index.
$limit and $skip: Restrict the number of documents returned or skip a number of documents. Often used for pagination.
$unwind: Deconstructs an array field into multiple documents, one per array element. This can increase the number of documents significantly.

Here's an example that combines several stages to find the total revenue per customer, sorted by total:

pipeline = [
    {'$match': {'status': 'completed'}},
    {'$unwind': '$items'},
    {'$group': {
        '_id': '$customer_id',
        'total': {'$sum': {'$multiply': ['$items.price', '$items.quantity']}}
    }},
    {'$sort': {'total': -1}},
    {'$limit': 10}
]

This pipeline first filters for completed orders, then expands each item in the items array, computes the line total, groups by customer, and sorts by the total. Notice how $unwind multiplies the number of documents; if an order has 5 items, it becomes 5 documents. This is a common pattern, but be aware of its performance implications on large collections.

Performance Considerations for Aggregation Pipelines

Performance in aggregation pipelines depends heavily on how you structure the stages and whether you use indexes. The most important rule is to filter as early as possible with $match. This reduces the number of documents that flow through the rest of the pipeline. For example, if you only need data from the last 30 days, a $match on a date field should be the first stage. Ensure that the field used in $match is indexed so MongoDB can use the index to quickly find matching documents.

Another performance factor is the use of $sort. If you sort after a $group, MongoDB must hold all the grouped results in memory before sorting. If the result set is large, this can cause high memory usage or even a server-side error. In such cases, consider using an index on the sort key if the sort is on a field that appears early in the pipeline. MongoDB can sometimes use an index to avoid an in-memory sort, but only if the index can satisfy the sort order and the previous stages preserve that order.

$unwind can also be expensive because it expands each array element into a separate document. If you only need to aggregate over array elements and do not need the individual elements later, you can sometimes use operators like $reduce or $sum directly on the array without unwinding. The correct expression depends on your data shape and MongoDB version, so test it with your actual data.

Working with Large Result Sets

When an aggregation returns many documents, the cursor behavior matters. By default, PyMongo fetches results in batches of 101 documents. You can control the batch size when constructing the cursor by passing batchSize to aggregate(). For example:

cursor = collection.aggregate(pipeline, batchSize=500)

This tells MongoDB to return up to 500 documents per batch, reducing the number of round trips to the server. However, a larger batch size also means more memory on the client side if you consume the entire cursor into a list. For very large results, iterate over the cursor rather than calling list(cursor) to avoid loading everything into memory at once.

You can also use the allowDiskUse option for aggregations that exceed the 100 MB memory limit for a stage. This allows MongoDB to spill data to disk. In PyMongo, you pass it as a parameter:

cursor = collection.aggregate(pipeline, allowDiskUse=True)

Use this option sparingly because disk-based sorting and grouping are slower than in-memory operations. It's better to restructure the pipeline to reduce memory usage, such as filtering earlier or using more selective $match stages.

Handling Errors in Aggregation Pipelines

Aggregation pipelines can fail for several reasons: invalid stage syntax, unsupported operators, or server-side memory limits. PyMongo raises pymongo.errors.OperationFailure when MongoDB returns an error. The error message usually contains details about the stage and the reason. For example, if you exceed the memory limit, you'll see a message like Exceeded memory limit for $group, but didn't allow external sort. To handle this, either add allowDiskUse=True or restructure the pipeline.

Another common source of incorrect results is a field name that doesn't exist on some documents. MongoDB doesn't throw an error for missing fields; it simply treats them as null. This can lead to unexpected results, especially with $group where null _id groups all documents together. To avoid this, validate your data or use $match to exclude documents without the field.

When debugging a pipeline, start by running it in the MongoDB shell with a small sample. You can also inspect the execution plan from PyMongo. The explain syntax varies by MongoDB version, but one common form is to send the aggregation command through db.command():

plan = db.command(
    'explain',
    {'aggregate': 'orders', 'pipeline': pipeline, 'cursor': {}},
    verbosity='executionStats',
)

The returned plan shows how MongoDB executes the pipeline, including whether indexes are used and how many documents are examined. This is invaluable for identifying performance bottlenecks. Keep in mind that the exact output shape depends on your MongoDB version.

Advanced Aggregation Patterns: $lookup and $facet

For more complex data transformations, you can use $lookup to perform a left outer join with another collection, and $facet to run multiple pipelines on the same input documents.

$lookup is useful when you need to combine data from two collections. For example, to attach customer details to each order:

pipeline = [
    {'$match': {'status': 'completed'}},
    {'$lookup': {
        'from': 'customers',
        'localField': 'customer_id',
        'foreignField': '_id',
        'as': 'customer'
    }},
    {'$unwind': '$customer'}
]

This adds a customer field containing the matching customer document. The $unwind is needed because $lookup always produces an array, even if there's only one match. Be careful with $lookup on large collections; the foreign field in the from collection should be indexed, otherwise the lookup can be slow.

$facet lets you compute multiple aggregations in a single pass over the input documents. For instance, you might want both the total count and the average order value:

pipeline = [
    {
        '$facet': {
            'total_orders': [{'$count': 'count'}],
            'avg_value': [
                {'$group': {'_id': None, 'avg': {'$avg': '$total'}}}
            ]
        }
    }
]

The output is a single document with two fields, each containing an array of results. This is efficient because MongoDB processes the input once and applies the sub-pipelines in the same stage. However, each sub-pipeline still has its own memory limit, so keep that in mind.

When combining $lookup and $facet, you can build powerful reporting queries that would otherwise require multiple round trips. Just remember to test with realistic data volumes to ensure the pipeline stays within MongoDB's memory and performance constraints.