24
24
25
25
namespace MongoDB . Driver . Linq . Linq3Implementation . Ast . Optimizers
26
26
{
27
- internal class AstGroupPipelineOptimizer
27
+ internal class AstGroupingPipelineOptimizer
28
28
{
29
29
#region static
30
30
public static AstPipeline Optimize ( AstPipeline pipeline )
31
31
{
32
- var optimizer = new AstGroupPipelineOptimizer ( ) ;
32
+ var optimizer = new AstGroupingPipelineOptimizer ( ) ;
33
33
for ( var i = 0 ; i < pipeline . Stages . Count ; i ++ )
34
34
{
35
35
var stage = pipeline . Stages [ i ] ;
36
- if ( stage is AstGroupStage groupStage )
36
+ if ( IsGroupingStage ( stage ) )
37
37
{
38
- pipeline = optimizer . OptimizeGroupStage ( pipeline , i , groupStage ) ;
38
+ pipeline = optimizer . OptimizeGroupingStage ( pipeline , i , stage ) ;
39
39
}
40
40
}
41
41
42
42
return pipeline ;
43
+
44
+ static bool IsGroupingStage ( AstStage stage )
45
+ {
46
+ return stage . NodeType switch
47
+ {
48
+ AstNodeType . GroupStage or AstNodeType . BucketStage or AstNodeType . BucketAutoStage => true ,
49
+ _ => false
50
+ } ;
51
+ }
43
52
}
44
53
#endregion
45
54
46
55
private readonly AccumulatorSet _accumulators = new AccumulatorSet ( ) ;
47
56
private AstExpression _element ; // normally either "$$ROOT" or "$_v"
48
57
49
- private AstPipeline OptimizeGroupStage ( AstPipeline pipeline , int i , AstGroupStage groupStage )
58
+ private AstPipeline OptimizeGroupingStage ( AstPipeline pipeline , int i , AstStage groupingStage )
50
59
{
51
60
try
52
61
{
53
- if ( IsOptimizableGroupStage ( groupStage , out _element ) )
62
+ if ( IsOptimizableGroupingStage ( groupingStage , out _element ) )
54
63
{
55
64
var followingStages = GetFollowingStagesToOptimize ( pipeline , i + 1 ) ;
56
65
if ( followingStages == null )
57
66
{
58
67
return pipeline ;
59
68
}
60
69
61
- var mappings = OptimizeGroupAndFollowingStages ( groupStage , followingStages ) ;
70
+ var mappings = OptimizeGroupingAndFollowingStages ( groupingStage , followingStages ) ;
62
71
if ( mappings . Length > 0 )
63
72
{
64
73
return ( AstPipeline ) AstNodeReplacer . Replace ( pipeline , mappings ) ;
@@ -72,23 +81,57 @@ private AstPipeline OptimizeGroupStage(AstPipeline pipeline, int i, AstGroupStag
72
81
73
82
return pipeline ;
74
83
75
- static bool IsOptimizableGroupStage ( AstGroupStage groupStage , out AstExpression element )
84
+ static bool IsOptimizableGroupingStage ( AstStage groupingStage , out AstExpression element )
76
85
{
77
- // { $group : { _id : ?, _elements : { $push : element } } }
78
- if ( groupStage . Fields . Count == 1 )
86
+ if ( groupingStage is AstGroupStage groupStage )
87
+ {
88
+ // { $group : { _id : ?, _elements : { $push : element } } }
89
+ if ( groupStage . Fields . Count == 1 )
90
+ {
91
+ var field = groupStage . Fields [ 0 ] ;
92
+ return IsElementsPush ( field , out element ) ;
93
+ }
94
+ }
95
+
96
+ if ( groupingStage is AstBucketStage bucketStage )
97
+ {
98
+ // { $bucket : { groupBy : ?, boundaries : ?, default : ?, output : { _elements : { $push : element } } } }
99
+ if ( bucketStage . Output . Count == 1 )
100
+ {
101
+ var output = bucketStage . Output [ 0 ] ;
102
+ return IsElementsPush ( output , out element ) ;
103
+ }
104
+ }
105
+
106
+ if ( groupingStage is AstBucketAutoStage bucketAutoStage )
79
107
{
80
- var field = groupStage . Fields [ 0 ] ;
81
- if ( field . Path == "_elements" &&
108
+ // { $bucketAuto : { groupBy : ?, buckets : ?, granularity : ?, output : { _elements : { $push : element } } } }
109
+ if ( bucketAutoStage . Output . Count == 1 )
110
+ {
111
+ var output = bucketAutoStage . Output [ 0 ] ;
112
+ return IsElementsPush ( output , out element ) ;
113
+ }
114
+ }
115
+
116
+ element = null ;
117
+ return false ;
118
+
119
+ static bool IsElementsPush ( AstAccumulatorField field , out AstExpression element )
120
+ {
121
+ if (
122
+ field . Path == "_elements" &&
82
123
field . Value is AstUnaryAccumulatorExpression unaryAccumulatorExpression &&
83
124
unaryAccumulatorExpression . Operator == AstUnaryAccumulatorOperator . Push )
84
125
{
85
126
element = unaryAccumulatorExpression . Arg ;
86
127
return true ;
87
128
}
129
+ else
130
+ {
131
+ element = null ;
132
+ return false ;
133
+ }
88
134
}
89
-
90
- element = null ;
91
- return false ;
92
135
}
93
136
94
137
static List < AstStage > GetFollowingStagesToOptimize ( AstPipeline pipeline , int from )
@@ -135,7 +178,7 @@ static bool IsLastStageThatCanBeOptimized(AstStage stage)
135
178
}
136
179
}
137
180
138
- private ( AstNode , AstNode ) [ ] OptimizeGroupAndFollowingStages ( AstGroupStage groupStage , List < AstStage > followingStages )
181
+ private ( AstNode , AstNode ) [ ] OptimizeGroupingAndFollowingStages ( AstStage groupingStage , List < AstStage > followingStages )
139
182
{
140
183
var mappings = new List < ( AstNode , AstNode ) > ( ) ;
141
184
@@ -148,10 +191,21 @@ static bool IsLastStageThatCanBeOptimized(AstStage stage)
148
191
}
149
192
}
150
193
151
- var newGroupStage = AstStage . Group ( groupStage . Id , _accumulators ) ;
152
- mappings . Add ( ( groupStage , newGroupStage ) ) ;
194
+ var newGroupingStage = CreateNewGroupingStage ( groupingStage , _accumulators ) ;
195
+ mappings . Add ( ( groupingStage , newGroupingStage ) ) ;
153
196
154
197
return mappings . ToArray ( ) ;
198
+
199
+ static AstStage CreateNewGroupingStage ( AstStage groupingStage , AccumulatorSet accumulators )
200
+ {
201
+ return groupingStage switch
202
+ {
203
+ AstGroupStage groupStage => AstStage . Group ( groupStage . Id , accumulators ) ,
204
+ AstBucketStage bucketStage => AstStage . Bucket ( bucketStage . GroupBy , bucketStage . Boundaries , bucketStage . Default , accumulators ) ,
205
+ AstBucketAutoStage bucketAutoStage => AstStage . BucketAuto ( bucketAutoStage . GroupBy , bucketAutoStage . Buckets , bucketAutoStage . Granularity , accumulators ) ,
206
+ _ => throw new Exception ( $ "Unexpected { nameof ( groupingStage ) } node type: { groupingStage . NodeType } .")
207
+ } ;
208
+ }
155
209
}
156
210
157
211
private AstStage OptimizeFollowingStage ( AstStage stage )
0 commit comments