MongoEngine查询返回所有匹配条件嵌入文档的实现方法

场景说明 #

基于Django + MongoEngine操作MongoDB数据时，投影阶段使用$elemMatch过滤数组类型的嵌入文档字段，仅能返回数组中第一个匹配条件的元素，无法满足返回所有匹配嵌入文档的业务需求。

模型定义 #

嵌入文档 Data #

class Data(EmbeddedDocument):
    v = FloatField()
    q = StringField()
    co2 = FloatField()
    price = FloatField()
    ts = DateTimeField()

主文档 MetersData #

class MetersData(DynamicDocument):
    meta = {'collection': 'dk_heating'}
    _id = ObjectIdField()
    ident = StringField()
    meteringPointId = StringField()
    customer = StringField()
    cvr = StringField()
    type = StringField()
    unit = StringField()
    address = StringField()
    period = EmbeddedDocumentField(Period)
    hourly_data = ListField(EmbeddedDocumentField(Data), db_field='data')
    daily_data = ListField(EmbeddedDocumentField(Data))
    monthly_data = ListField(EmbeddedDocumentField(Data))
    # monthly_data = EmbeddedDocumentListField(Data)
    yearly_data = ListField(EmbeddedDocumentField(Data))

原有问题查询 #

原查询写法如下，该写法受$elemMatch投影的设计限制，仅能返回monthly_data数组中第一个q="E"的元素：

MetersData.objects.filter(address=address, customer=customer).fields(
                monthly_data={"$elemMatch": {"q": "E"}},
                address=1, customer=1, cvr=1, ident=1, meteringPointId=1, type=1, unit=1, period=1)

可行实现方案 #

方案1：聚合管道$filter实现（推荐，MongoDB 3.2+版本支持） #

使用聚合管道的$filter操作符可以在数据库层面完成数组过滤，直接返回所有符合条件的数组元素，性能开销最低，MongoEngine可直接调用聚合接口实现：

pipeline = [
    # 匹配主文档层面的过滤条件，与原查询filter逻辑一致
    {"$match": {"address": address, "customer": customer}},
    # 指定返回字段，同时对目标数组字段做全量条件过滤
    {"$project": {
        "address": 1,
        "customer": 1,
        "cvr": 1,
        "ident": 1,
        "meteringPointId": 1,
        "type": 1,
        "unit": 1,
        "period": 1,
        # 过滤monthly_data数组，保留所有q="E"的元素
        "monthly_data": {
            "$filter": {
                "input": "$monthly_data",
                "as": "item",
                "cond": {"$eq": ["$$item.q", "E"]}
            }
        }
    }}
]
# 执行聚合查询获取结果
result = MetersData.objects.aggregate(pipeline)

如果需要同时过滤hourly_data/daily_data等其他数组字段，在$project块中按相同语法添加对应字段的过滤规则即可。 #

方案2：字段类型替换+代码层面二次过滤 #

如果不想使用聚合管道，可以将monthly_data的字段类型从普通ListField(EmbeddedDocumentField(Data))替换为代码中注释的EmbeddedDocumentListField(Data)，该类型是MongoEngine专为嵌入文档数组设计的字段类型，支持数组元素的条件匹配：

# 第一步：修改字段定义
monthly_data = EmbeddedDocumentListField(Data)

# 第二步：查询匹配包含目标元素的主文档
queryset = MetersData.objects.filter(
    address=address, 
    customer=customer,
    monthly_data__q="E"
).only("address", "customer", "cvr", "ident", "meteringPointId", "type", "unit", "period", "monthly_data")

# 第三步：代码层面二次过滤每个文档的monthly_data数组，保留符合条件的元素
for doc in queryset:
    doc.monthly_data = [item for item in doc.monthly_data if item.q == "E"]
    # 后续直接使用处理后的doc即可

注意：该方案会把匹配主文档的全量monthly_data数组加载到内存后再过滤，数组长度较大时网络传输和内存开销明显高于方案1，非特殊场景优先选择方案1。

内容的提问来源于stack exchange，提问作者Muhammad Haris Yousaf