43
43
import com .apple .foundationdb .record .query .plan .cascades .typing .TypeRepository ;
44
44
import com .apple .foundationdb .record .query .plan .cascades .values .FieldValue ;
45
45
import com .apple .foundationdb .record .query .plan .cascades .values .QuantifiedObjectValue ;
46
+ import com .apple .foundationdb .record .query .plan .cascades .values .RecordConstructorValue ;
46
47
import com .apple .foundationdb .record .query .plan .cascades .values .Value ;
47
48
import com .apple .foundationdb .record .query .plan .cascades .values .Values ;
48
49
import com .apple .foundationdb .record .query .plan .cascades .values .simplification .OrderingValueComputationRuleSet ;
50
+ import com .apple .foundationdb .record .query .plan .cascades .values .translation .PullUp ;
49
51
import com .apple .foundationdb .record .query .plan .plans .RecordQueryAggregateIndexPlan ;
50
52
import com .apple .foundationdb .record .query .plan .plans .RecordQueryFetchFromPartialRecordPlan ;
51
53
import com .apple .foundationdb .record .query .plan .plans .RecordQueryIndexPlan ;
52
54
import com .apple .foundationdb .record .query .plan .plans .RecordQueryPlan ;
55
+ import com .apple .foundationdb .record .query .plan .plans .RecordQueryStreamingAggregationPlan ;
53
56
import com .apple .foundationdb .record .util .pair .NonnullPair ;
54
57
import com .google .common .base .Preconditions ;
55
58
import com .google .common .base .Verify ;
@@ -172,7 +175,7 @@ public KeyExpression getFullKeyExpression() {
172
175
173
176
@ Override
174
177
public String toString () {
175
- return "Agg [" + getName () + "; " + index .getType () + "]" ;
178
+ return "AGG [" + getName () + "; " + index .getType () + "]" ;
176
179
}
177
180
178
181
@ Override
@@ -374,6 +377,33 @@ public Ordering computeOrderingFromScanComparisons(@Nonnull final ScanComparison
374
377
return Ordering .ofOrderingSequence (bindingMapBuilder .build (), orderingSequenceBuilder .build (), isDistinct );
375
378
}
376
379
380
+ @ Nullable
381
+ @ Override
382
+ public PullUp .UnificationPullUp prepareForUnification (@ Nonnull final PartialMatch partialMatch ,
383
+ @ Nonnull final CorrelationIdentifier topAlias ,
384
+ @ Nonnull final CorrelationIdentifier topCandidateAlias ) {
385
+ final var regularMatchInfo = partialMatch .getRegularMatchInfo ();
386
+ if (regularMatchInfo .getRollUpToGroupingValues () != null ) {
387
+ final var groupingAndAggregateAccessors =
388
+ getGroupingAndAggregateAccessors (topCandidateAlias );
389
+ final var groupingAccessorValues = groupingAndAggregateAccessors .getLeft ();
390
+ final var aggregateAccessorValue = groupingAndAggregateAccessors .getRight ();
391
+ final var allFields =
392
+ ((Type .Record )selectHavingExpression .getResultValue ().getResultType ()).getFields ();
393
+ final var rollUpColumnsBuilder =
394
+ ImmutableList .<Column <? extends Value >>builder ();
395
+ final var numGroupings = regularMatchInfo .getRollUpToGroupingValues ().size ();
396
+ for (int i = 0 ; i < numGroupings ; i ++) {
397
+ final var field = allFields .get (i );
398
+ rollUpColumnsBuilder .add (Column .of (field , groupingAccessorValues .get (i )));
399
+ }
400
+ rollUpColumnsBuilder .add (Column .of (allFields .get (allFields .size () - 1 ), aggregateAccessorValue ));
401
+ return PullUp .forUnification (topAlias , RecordConstructorValue .ofColumns (rollUpColumnsBuilder .build ()),
402
+ ImmutableSet .of (topCandidateAlias ));
403
+ }
404
+ return null ;
405
+ }
406
+
377
407
@ Nonnull
378
408
@ Override
379
409
public RecordQueryPlan toEquivalentPlan (@ Nonnull final PartialMatch partialMatch ,
@@ -384,13 +414,14 @@ public RecordQueryPlan toEquivalentPlan(@Nonnull final PartialMatch partialMatch
384
414
final var baseRecordType = Type .Record .fromFieldDescriptorsMap (RecordMetaData .getFieldDescriptorMapFromTypes (recordTypes ));
385
415
386
416
final var selectHavingResultValue = selectHavingExpression .getResultValue ();
387
- final var resultType = selectHavingResultValue .getResultType ();
417
+ final var resultType = ( Type . Record ) selectHavingResultValue .getResultType ();
388
418
final var messageDescriptor =
389
419
Objects .requireNonNull (TypeRepository .newBuilder ()
390
420
.addTypeIfNeeded (resultType )
391
421
.build ()
392
422
.getMessageDescriptor (resultType ));
393
- final var constraintMaybe = partialMatch .getRegularMatchInfo ().getConstraint ();
423
+ final var regularMatchInfo = partialMatch .getRegularMatchInfo ();
424
+ final var constraintMaybe = regularMatchInfo .getConstraint ();
394
425
395
426
final var indexEntryConverter = createIndexEntryConverter (messageDescriptor );
396
427
final var aggregateIndexScan = new RecordQueryIndexPlan (index .getName (),
@@ -404,12 +435,45 @@ public RecordQueryPlan toEquivalentPlan(@Nonnull final PartialMatch partialMatch
404
435
baseRecordType ,
405
436
QueryPlanConstraint .noConstraint ());
406
437
407
- return new RecordQueryAggregateIndexPlan (aggregateIndexScan ,
438
+ var plan = new RecordQueryAggregateIndexPlan (aggregateIndexScan ,
408
439
recordTypes .get (0 ).getName (),
409
440
indexEntryConverter ,
410
441
selectHavingResultValue ,
411
442
groupByResultValue ,
412
443
constraintMaybe );
444
+ if (regularMatchInfo .getRollUpToGroupingValues () != null ) {
445
+ //
446
+ // We need to perform a roll up.
447
+ //
448
+ final var aggregateIndexScanReference = memoizer .memoizePlans (plan );
449
+ final var aggregateIndexScanAlias = Quantifier .uniqueId ();
450
+
451
+ //final var recordValues = Values.deconstructRecord(recordValue);
452
+ final var groupingAndAggregateAccessors =
453
+ getGroupingAndAggregateAccessors (regularMatchInfo .getRollUpToGroupingValues ().size (),
454
+ aggregateIndexScanAlias );
455
+ final var groupingAccessorValues = groupingAndAggregateAccessors .getLeft ();
456
+ final var aggregateAccessorValue = groupingAndAggregateAccessors .getRight ();
457
+ final var allFields = resultType .getFields ();
458
+ final var rollUpGroupingColumnsBuilder =
459
+ ImmutableList .<Column <? extends Value >>builder ();
460
+ for (int i = 0 ; i < groupingAccessorValues .size (); i ++) {
461
+ final var field = allFields .get (i );
462
+ rollUpGroupingColumnsBuilder .add (Column .of (field , groupingAccessorValues .get (i )));
463
+ }
464
+
465
+ final var rollUpAggregateValueOptional =
466
+ AggregateIndexExpansionVisitor .rollUpAggregateValueMaybe (index ,
467
+ aggregateAccessorValue );
468
+
469
+ final var aggregateIndexScanQuantifier =
470
+ Quantifier .physical (aggregateIndexScanReference , aggregateIndexScanAlias );
471
+
472
+ return RecordQueryStreamingAggregationPlan .ofFlattened (aggregateIndexScanQuantifier ,
473
+ RecordConstructorValue .ofColumns (rollUpGroupingColumnsBuilder .build (), resultType .isNullable ()),
474
+ rollUpAggregateValueOptional .orElseThrow (() -> new RecordCoreException ("unknown rollup operation" )));
475
+ }
476
+ return plan ;
413
477
}
414
478
415
479
@ Nonnull
@@ -426,19 +490,22 @@ protected int getGroupingCount() {
426
490
}
427
491
428
492
@ Nonnull
429
- public NonnullPair <List <Value >, Value > getGroupingAndAggregateAccessors () {
493
+ public NonnullPair <List <Value >, Value > getGroupingAndAggregateAccessors (@ Nonnull final CorrelationIdentifier alias ) {
494
+ return getGroupingAndAggregateAccessors (getGroupingCount (), alias );
495
+ }
496
+
497
+ @ Nonnull
498
+ public NonnullPair <List <Value >, Value > getGroupingAndAggregateAccessors (final int numGroupings ,
499
+ @ Nonnull final CorrelationIdentifier alias ) {
430
500
final var selectHavingResultValue = selectHavingExpression .getResultValue ();
431
501
final var selectHavingResultType = (Type .Record )selectHavingResultValue .getResultType ();
432
502
final var unbasedResultValue =
433
- QuantifiedObjectValue .of (Quantifier .current (), selectHavingResultType );
434
-
435
- final var groupingCount = getGroupingCount ();
436
- Verify .verify (selectHavingResultType .getFields ().size () >= groupingCount );
437
-
503
+ QuantifiedObjectValue .of (alias , selectHavingResultType );
438
504
final var deconstructedRecord = Values .deconstructRecord (unbasedResultValue );
505
+ Verify .verify (deconstructedRecord .size () > numGroupings );
439
506
final var groupingAccessorValues =
440
- ImmutableList .copyOf (deconstructedRecord .subList (0 , groupingCount ));
441
- final var aggregateAccessorValue = deconstructedRecord .get (groupingCount );
507
+ ImmutableList .copyOf (deconstructedRecord .subList (0 , numGroupings ));
508
+ final var aggregateAccessorValue = deconstructedRecord .get (deconstructedRecord . size () - 1 );
442
509
return NonnullPair .of (groupingAccessorValues , aggregateAccessorValue );
443
510
}
444
511
0 commit comments