Skip to main content

vantage_mongodb/select/
pipeline.rs

1//! Aggregation pipeline builders for MongoSelect. Every pipeline starts with
2//! `pre_stages`, then the filter as `$match`. Join stages run as late as the
3//! read allows: after the conditions on stored fields, and after paging too
4//! unless a condition or the sort reads a joined column.
5
6use bson::{Bson, Document, doc};
7
8use super::MongoSelect;
9use super::join::filter_uses_joins;
10use crate::condition::merge_documents;
11
12impl MongoSelect {
13    /// The filter split into its conditions on stored fields and those
14    /// reading a joined column.
15    async fn split_filter(&self) -> vantage_core::Result<(Document, Document)> {
16        let (mut stored, mut joined) = (Vec::new(), Vec::new());
17        for condition in &self.conditions {
18            let resolved = condition.resolve().await?;
19            if filter_uses_joins(&resolved, self.joins.len()) {
20                joined.push(resolved);
21            } else {
22                stored.push(resolved);
23            }
24        }
25        Ok((merge_documents(stored)?, merge_documents(joined)?))
26    }
27
28    /// The filter as `$match` stages: one on stored fields, one on joined
29    /// columns; `None` where there are no such conditions.
30    async fn match_stages(&self) -> vantage_core::Result<(Option<Document>, Option<Document>)> {
31        let (stored, joined) = self.split_filter().await?;
32        let stage = |filter: Document| (!filter.is_empty()).then(|| doc! { "$match": filter });
33        Ok((stage(stored), stage(joined)))
34    }
35
36    /// `pre_stages` and the stored-field `$match`; then, when a condition
37    /// reads a joined column, the join stages and the joined `$match`.
38    async fn filtered_stages(&self) -> vantage_core::Result<Vec<Document>> {
39        let (stored, joined) = self.match_stages().await?;
40        let mut pipeline = self.pre_stages.clone();
41        pipeline.extend(stored);
42        if let Some(joined) = joined {
43            pipeline.extend(self.join_stages());
44            pipeline.push(joined);
45        }
46        Ok(pipeline)
47    }
48
49    /// The layout of a row read, given its two `$match` stages in whatever
50    /// form `stage` renders documents to:
51    /// [pre_stages.., stored $match, $sort, $skip, $limit, joins.., $project],
52    /// with the joins (and the joined `$match`) moved ahead of the paging when
53    /// a condition or the sort reads a joined column.
54    pub(super) fn read_layout<S>(
55        &self,
56        stored: Option<S>,
57        joined: Option<S>,
58        stage: impl Fn(Document) -> S,
59    ) -> Vec<S> {
60        let mut out: Vec<S> = self.pre_stages.iter().cloned().map(&stage).collect();
61        out.extend(stored);
62        let joins = self.join_stages().into_iter().map(&stage);
63        let page = self.page_stages().into_iter().map(&stage);
64        if joined.is_some() || self.sort_uses_joins() {
65            out.extend(joins);
66            out.extend(joined);
67            out.extend(page);
68        } else {
69            out.extend(page);
70            out.extend(joins);
71        }
72        out.extend(self.projection_stage().map(&stage));
73        out
74    }
75
76    /// [$sort, $skip, $limit], each only when set.
77    fn page_stages(&self) -> Vec<Document> {
78        let mut stages = Vec::new();
79        if let Some(sort) = self.build_sort() {
80            stages.push(doc! { "$sort": sort });
81        }
82        if let Some(skip) = self.skip.filter(|s| *s > 0) {
83            stages.push(doc! { "$skip": skip });
84        }
85        if let Some(limit) = self.limit.filter(|l| *l > 0) {
86            stages.push(doc! { "$limit": limit });
87        }
88        stages
89    }
90
91    fn projection_stage(&self) -> Option<Document> {
92        self.build_projection()
93            .map(|projection| doc! { "$project": projection })
94    }
95
96    /// Build a row-reading pipeline; see `read_layout` for the stage order.
97    pub async fn as_read_pipeline(&self) -> vantage_core::Result<Vec<Document>> {
98        let (stored, joined) = self.match_stages().await?;
99        Ok(self.read_layout(stored, joined, |d| d))
100    }
101
102    /// Build a count aggregation pipeline: [pre_stages.., $match, $count],
103    /// with the join stages only when the filter reads a joined column.
104    pub async fn as_count_pipeline(&self) -> vantage_core::Result<Vec<Document>> {
105        let mut pipeline = self.filtered_stages().await?;
106        pipeline.push(doc! { "$count": "count" });
107        Ok(pipeline)
108    }
109
110    /// Build a $group aggregation pipeline with an accumulator.
111    /// `func` is e.g. "$sum", "$max", "$min". `field` is the column name.
112    pub async fn as_aggregate_pipeline(
113        &self,
114        func: &str,
115        field: &str,
116    ) -> vantage_core::Result<Vec<Document>> {
117        let field_ref = format!("${}", field);
118        let mut pipeline = self.filtered_stages().await?;
119        pipeline.push(doc! { "$group": { "_id": Bson::Null, "val": { func: field_ref } } });
120        Ok(pipeline)
121    }
122}