python pymongo aggregation pipeline: Practical Usage and Cod

When you need to transform, filter, and aggregate documents in MongoDB from a Python application, the aggregation pipeline is the tool to use. In this article, we'll build a python pymongo aggregation pipeline from scratch, covering the essential stages, performance considerations, and common pitfalls. You'll learn how to construct pipelines that run efficiently and produce 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 stage operators as the MongoDB shell. 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 an 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 without needing individual elements later, consider using aggregation operators like $sum with an array expression directly. For example, instead of unwinding and then grouping, you can use $sum with a mapped array in a $project stage, but this requires MongoDB 3.4 or later. Always test with your actual data volume.

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 error is using a field name that doesn't exist. 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 use the explain command in PyMongo to see how MongoDB executes the pipeline and whether indexes are used:

result = collection.aggregate(pipeline, explain=True)

The output shows the execution plan, including which stages used indexes and how many documents were examined. This is invaluable for identifying performance bottlenecks.

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 in parallel 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; it can be slow if the foreign field isn't indexed.

$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": null, "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 runs the sub-pipelines in parallel. 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.