34
34
import com .apple .foundationdb .record .query .plan .cascades .expressions .SelectExpression ;
35
35
import com .apple .foundationdb .record .query .plan .cascades .predicates .Placeholder ;
36
36
import com .apple .foundationdb .record .query .plan .cascades .predicates .PredicateWithValueAndRanges ;
37
+ import com .apple .foundationdb .record .query .plan .cascades .values .AggregateValue ;
37
38
import com .apple .foundationdb .record .query .plan .cascades .values .ArithmeticValue ;
38
39
import com .apple .foundationdb .record .query .plan .cascades .values .CountValue ;
39
40
import com .apple .foundationdb .record .query .plan .cascades .values .EmptyValue ;
57
58
import javax .annotation .Nonnull ;
58
59
import javax .annotation .Nullable ;
59
60
import java .util .Collection ;
60
- import java .util .Collections ;
61
61
import java .util .List ;
62
62
import java .util .Map ;
63
63
import java .util .Objects ;
64
+ import java .util .Optional ;
64
65
import java .util .function .Supplier ;
65
66
import java .util .stream .Stream ;
66
67
72
73
public class AggregateIndexExpansionVisitor extends KeyExpressionExpansionVisitor
73
74
implements ExpansionVisitor <KeyExpressionExpansionVisitor .VisitorState > {
74
75
@ Nonnull
75
- static final Supplier <Map <String , BuiltInFunction <? extends Value >>> aggregateMap = Suppliers .memoize (AggregateIndexExpansionVisitor ::computeAggregateMap );
76
+ static final Supplier <Map <String , BuiltInFunction <? extends Value >>> aggregateMap =
77
+ Suppliers .memoize (AggregateIndexExpansionVisitor ::computeAggregateMap );
78
+
79
+ @ Nonnull
80
+ static final Supplier <Map <String , BuiltInFunction <? extends Value >>> rollUpAggregateMap =
81
+ Suppliers .memoize (AggregateIndexExpansionVisitor ::computeRollUpAggregateMap );
76
82
77
83
@ Nonnull
78
84
protected final Index index ;
@@ -92,7 +98,8 @@ public class AggregateIndexExpansionVisitor extends KeyExpressionExpansionVisito
92
98
* @param recordTypes The indexed record types.
93
99
*/
94
100
public AggregateIndexExpansionVisitor (@ Nonnull final Index index , @ Nonnull final Collection <RecordType > recordTypes ) {
95
- Preconditions .checkArgument (IndexTypes .BITMAP_VALUE .equals (index .getType ()) || aggregateMap .get ().containsKey (index .getType ()));
101
+ Preconditions .checkArgument (IndexTypes .BITMAP_VALUE .equals (index .getType ()) ||
102
+ aggregateMap .get ().containsKey (index .getType ()));
96
103
Preconditions .checkArgument (index .getRootExpression () instanceof GroupingKeyExpression );
97
104
this .index = index ;
98
105
this .groupingKeyExpression = ((GroupingKeyExpression )index .getRootExpression ());
@@ -139,9 +146,9 @@ public MatchCandidate expand(@Nonnull final Supplier<Quantifier.ForEach> baseQua
139
146
.addAll (groupByPlaceholders ).build ();
140
147
141
148
// 3. construct SELECT-HAVING with SORT on top.
142
- final var selectHavingAndPlaceholderAliases = constructSelectHaving (groupByQun , placeholders );
143
- final var selectHaving = selectHavingAndPlaceholderAliases . getLeft ();
144
- final var placeHolderAliases = selectHavingAndPlaceholderAliases . getRight ();
149
+ final var constructSelectHavingResult = constructSelectHaving (groupByQun , placeholders );
150
+ final var selectHaving = constructSelectHavingResult . getSelectExpression ();
151
+ final var placeHolderAliases = constructSelectHavingResult . getPlaceholderAliases ();
145
152
146
153
// 4. add sort on top, if necessary, this will be absorbed later on as an ordering property of the match candidate.
147
154
final var maybeWithSort = placeHolderAliases .isEmpty ()
@@ -243,7 +250,8 @@ protected NonnullPair<Quantifier, List<Placeholder>> constructGroupBy(@Nonnull f
243
250
throw new RecordCoreException ("unable to plan group by with non-field value" )
244
251
.addLogInfo (LogMessageKeys .VALUE , groupedValue );
245
252
}
246
- final var aggregateValue = (Value )aggregateMap .get ().get (index .getType ()).encapsulate (ImmutableList .of (argument ));
253
+ final var aggregateValue =
254
+ aggregateValue (index , argument ).orElseThrow (() -> new RecordCoreException ("unknown aggregation type" ));
247
255
// add an RCV column representing the grouping columns as the first result set column
248
256
// also, make sure to set the field type names correctly for each field value in the grouping keys RCV.
249
257
@@ -278,19 +286,25 @@ protected NonnullPair<Quantifier, List<Placeholder>> constructGroupBy(@Nonnull f
278
286
}
279
287
280
288
@ Nonnull
281
- private NonnullPair <SelectExpression , List <CorrelationIdentifier >> constructSelectHaving (@ Nonnull final Quantifier groupByQun ,
282
- @ Nonnull final List <Placeholder > selectWherePlaceholders ) {
289
+ private ConstructSelectHavingResult constructSelectHaving (@ Nonnull final Quantifier groupByQun ,
290
+ @ Nonnull final List <Placeholder > selectWherePlaceholders ) {
291
+ final var rangesOverExpression = groupByQun .getRangesOver ().get ();
292
+ Verify .verify (rangesOverExpression instanceof GroupByExpression );
293
+ final var groupByExpression = (GroupByExpression )rangesOverExpression ;
294
+
283
295
// the grouping value in GroupByExpression comes first (if set).
284
296
@ Nullable final var groupingValueReference =
285
- (groupByQun .getRangesOver ().get () instanceof GroupByExpression && ((GroupByExpression )groupByQun .getRangesOver ().get ()).getGroupingValue () == null )
286
- ? null
287
- : FieldValue .ofOrdinalNumber (groupByQun .getFlowedObjectValue (), 0 );
297
+ groupByExpression .getGroupingValue () == null
298
+ ? null : FieldValue .ofOrdinalNumber (groupByQun .getFlowedObjectValue (), 0 );
288
299
289
- final var aggregateValueReference = FieldValue .ofOrdinalNumberAndFuseIfPossible (FieldValue .ofOrdinalNumber (groupByQun .getFlowedObjectValue (), groupingValueReference == null ? 0 : 1 ), 0 );
300
+ final var aggregateValueReference =
301
+ FieldValue .ofOrdinalNumberAndFuseIfPossible (FieldValue .ofOrdinalNumber (groupByQun .getFlowedObjectValue (),
302
+ groupingValueReference == null ? 0 : 1 ), 0 );
290
303
291
304
final var placeholderAliases = ImmutableList .<CorrelationIdentifier >builder ();
292
305
final var selectHavingGraphExpansionBuilder = GraphExpansion .builder ().addQuantifier (groupByQun );
293
- final List <Value > groupingValues = groupingValueReference == null ? Collections .emptyList () : Values .deconstructRecord (groupingValueReference );
306
+ final List <Value > groupingValues = groupingValueReference == null
307
+ ? ImmutableList .of () : Values .deconstructRecord (groupingValueReference );
294
308
if (groupingValueReference != null ) {
295
309
int i = 0 ;
296
310
for (final var groupingValue : groupingValues ) {
@@ -322,7 +336,15 @@ private NonnullPair<SelectExpression, List<CorrelationIdentifier>> constructSele
322
336
} else {
323
337
finalPlaceholders = placeholderAliases .build ();
324
338
}
325
- return NonnullPair .of (selectHavingGraphExpansionBuilder .build ().buildSelect (), finalPlaceholders );
339
+
340
+ return new ConstructSelectHavingResult (selectHavingGraphExpansionBuilder .build ().buildSelect (),
341
+ finalPlaceholders );
342
+ }
343
+
344
+ @ Nonnull
345
+ public static Optional <AggregateValue > aggregateValue (@ Nonnull final Index index , @ Nonnull final Value argument ) {
346
+ return Optional .of ((AggregateValue )aggregateMap .get ()
347
+ .get (index .getType ()).encapsulate (ImmutableList .of (argument )));
326
348
}
327
349
328
350
@ Nonnull
@@ -339,4 +361,53 @@ private static Map<String, BuiltInFunction<? extends Value>> computeAggregateMap
339
361
mapBuilder .put (IndexTypes .PERMUTED_MIN , new NumericAggregationValue .MinFn ());
340
362
return mapBuilder .build ();
341
363
}
364
+
365
+ public static boolean canBeRolledUp (@ Nonnull final String indexType ) {
366
+ return rollUpAggregateMap .get ().containsKey (indexType );
367
+ }
368
+
369
+ @ Nonnull
370
+ public static Optional <AggregateValue > rollUpAggregateValueMaybe (@ Nonnull final String indexType , @ Nonnull final Value argument ) {
371
+ return Optional .ofNullable (rollUpAggregateMap .get ()
372
+ .get (indexType ))
373
+ .map (fn -> (AggregateValue )fn .encapsulate (ImmutableList .of (argument )));
374
+ }
375
+
376
+ @ Nonnull
377
+ private static Map <String , BuiltInFunction <? extends Value >> computeRollUpAggregateMap () {
378
+ final ImmutableMap .Builder <String , BuiltInFunction <? extends Value >> mapBuilder = ImmutableMap .builder ();
379
+ mapBuilder .put (IndexTypes .MAX_EVER_LONG , new NumericAggregationValue .MaxFn ());
380
+ mapBuilder .put (IndexTypes .MIN_EVER_LONG , new NumericAggregationValue .MinFn ());
381
+ mapBuilder .put (IndexTypes .MAX_EVER_TUPLE , new NumericAggregationValue .MaxFn ());
382
+ mapBuilder .put (IndexTypes .MIN_EVER_TUPLE , new NumericAggregationValue .MinFn ());
383
+ mapBuilder .put (IndexTypes .SUM , new NumericAggregationValue .SumFn ());
384
+ mapBuilder .put (IndexTypes .COUNT , new NumericAggregationValue .SumFn ());
385
+ mapBuilder .put (IndexTypes .COUNT_NOT_NULL , new NumericAggregationValue .SumFn ());
386
+ mapBuilder .put (IndexTypes .PERMUTED_MAX , new NumericAggregationValue .MaxFn ());
387
+ mapBuilder .put (IndexTypes .PERMUTED_MIN , new NumericAggregationValue .MinFn ());
388
+ return mapBuilder .build ();
389
+ }
390
+
391
+ private static class ConstructSelectHavingResult {
392
+ @ Nonnull
393
+ private final SelectExpression selectExpression ;
394
+ @ Nonnull
395
+ private final List <CorrelationIdentifier > placeholderAliases ;
396
+
397
+ private ConstructSelectHavingResult (@ Nonnull final SelectExpression selectExpression ,
398
+ @ Nonnull final List <CorrelationIdentifier > placeholderAliases ) {
399
+ this .selectExpression = selectExpression ;
400
+ this .placeholderAliases = placeholderAliases ;
401
+ }
402
+
403
+ @ Nonnull
404
+ public SelectExpression getSelectExpression () {
405
+ return selectExpression ;
406
+ }
407
+
408
+ @ Nonnull
409
+ public List <CorrelationIdentifier > getPlaceholderAliases () {
410
+ return placeholderAliases ;
411
+ }
412
+ }
342
413
}
0 commit comments