Skip to content

Commit aa05b0f

Browse files
normen662normen662
committed
roll ups work
1 parent 4c9ec49 commit aa05b0f

20 files changed

+489
-146
lines changed

‎fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/AggregateIndexExpansionVisitor.java

Lines changed: 38 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
import com.apple.foundationdb.record.query.plan.cascades.expressions.SelectExpression;
3434
import com.apple.foundationdb.record.query.plan.cascades.predicates.Placeholder;
3535
import com.apple.foundationdb.record.query.plan.cascades.predicates.PredicateWithValueAndRanges;
36+
import com.apple.foundationdb.record.query.plan.cascades.values.AggregateValue;
3637
import com.apple.foundationdb.record.query.plan.cascades.values.ArithmeticValue;
3738
import com.apple.foundationdb.record.query.plan.cascades.values.CountValue;
3839
import com.apple.foundationdb.record.query.plan.cascades.values.EmptyValue;
@@ -60,6 +61,7 @@
6061
import java.util.List;
6162
import java.util.Map;
6263
import java.util.Objects;
64+
import java.util.Optional;
6365
import java.util.function.Supplier;
6466
import java.util.stream.Stream;
6567

@@ -71,7 +73,12 @@
7173
public class AggregateIndexExpansionVisitor extends KeyExpressionExpansionVisitor
7274
implements ExpansionVisitor<KeyExpressionExpansionVisitor.VisitorState> {
7375
@Nonnull
74-
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);
7582

7683
@Nonnull
7784
protected final Index index;
@@ -91,7 +98,8 @@ public class AggregateIndexExpansionVisitor extends KeyExpressionExpansionVisito
9198
* @param recordTypes The indexed record types.
9299
*/
93100
public AggregateIndexExpansionVisitor(@Nonnull final Index index, @Nonnull final Collection<RecordType> recordTypes) {
94-
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()));
95103
Preconditions.checkArgument(index.getRootExpression() instanceof GroupingKeyExpression);
96104
this.index = index;
97105
this.groupingKeyExpression = ((GroupingKeyExpression)index.getRootExpression());
@@ -334,6 +342,12 @@ private ConstructSelectHavingResult constructSelectHaving(@Nonnull final Quantif
334342
finalPlaceholders, currentGroupingValues);
335343
}
336344

345+
@Nonnull
346+
public static Optional<AggregateValue> aggregateValue(@Nonnull final Index index, @Nonnull final Value argument) {
347+
return Optional.of((AggregateValue)aggregateMap.get()
348+
.get(index.getType()).encapsulate(ImmutableList.of(argument)));
349+
}
350+
337351
@Nonnull
338352
private static Map<String, BuiltInFunction<? extends Value>> computeAggregateMap() {
339353
final ImmutableMap.Builder<String, BuiltInFunction<? extends Value>> mapBuilder = ImmutableMap.builder();
@@ -349,6 +363,28 @@ private static Map<String, BuiltInFunction<? extends Value>> computeAggregateMap
349363
return mapBuilder.build();
350364
}
351365

366+
@Nonnull
367+
public static Optional<AggregateValue> rollUpAggregateValueMaybe(@Nonnull final Index index, @Nonnull final Value argument) {
368+
return Optional.ofNullable(rollUpAggregateMap.get()
369+
.get(index.getType()))
370+
.map(fn -> (AggregateValue)fn.encapsulate(ImmutableList.of(argument)));
371+
}
372+
373+
@Nonnull
374+
private static Map<String, BuiltInFunction<? extends Value>> computeRollUpAggregateMap() {
375+
final ImmutableMap.Builder<String, BuiltInFunction<? extends Value>> mapBuilder = ImmutableMap.builder();
376+
mapBuilder.put(IndexTypes.MAX_EVER_LONG, new NumericAggregationValue.MaxFn());
377+
mapBuilder.put(IndexTypes.MIN_EVER_LONG, new NumericAggregationValue.MinFn());
378+
// mapBuilder.put(IndexTypes.MAX_EVER_TUPLE, TODO);
379+
// mapBuilder.put(IndexTypes.MIN_EVER_TUPLE, TODO);
380+
mapBuilder.put(IndexTypes.SUM, new NumericAggregationValue.SumFn());
381+
mapBuilder.put(IndexTypes.COUNT, new NumericAggregationValue.SumFn());
382+
mapBuilder.put(IndexTypes.COUNT_NOT_NULL, new NumericAggregationValue.SumFn());
383+
mapBuilder.put(IndexTypes.PERMUTED_MAX, new NumericAggregationValue.MaxFn());
384+
mapBuilder.put(IndexTypes.PERMUTED_MIN, new NumericAggregationValue.MinFn());
385+
return mapBuilder.build();
386+
}
387+
352388
private static class ConstructSelectHavingResult {
353389
@Nonnull
354390
private final SelectExpression selectExpression;

‎fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/AggregateIndexMatchCandidate.java

Lines changed: 79 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -42,13 +42,16 @@
4242
import com.apple.foundationdb.record.query.plan.cascades.typing.TypeRepository;
4343
import com.apple.foundationdb.record.query.plan.cascades.values.FieldValue;
4444
import com.apple.foundationdb.record.query.plan.cascades.values.QuantifiedObjectValue;
45+
import com.apple.foundationdb.record.query.plan.cascades.values.RecordConstructorValue;
4546
import com.apple.foundationdb.record.query.plan.cascades.values.Value;
4647
import com.apple.foundationdb.record.query.plan.cascades.values.Values;
4748
import com.apple.foundationdb.record.query.plan.cascades.values.simplification.OrderingValueComputationRuleSet;
49+
import com.apple.foundationdb.record.query.plan.cascades.values.translation.PullUp;
4850
import com.apple.foundationdb.record.query.plan.plans.RecordQueryAggregateIndexPlan;
4951
import com.apple.foundationdb.record.query.plan.plans.RecordQueryFetchFromPartialRecordPlan;
5052
import com.apple.foundationdb.record.query.plan.plans.RecordQueryIndexPlan;
5153
import com.apple.foundationdb.record.query.plan.plans.RecordQueryPlan;
54+
import com.apple.foundationdb.record.query.plan.plans.RecordQueryStreamingAggregationPlan;
5255
import com.apple.foundationdb.record.util.pair.NonnullPair;
5356
import com.google.common.base.Preconditions;
5457
import com.google.common.base.Verify;
@@ -171,7 +174,7 @@ public KeyExpression getFullKeyExpression() {
171174

172175
@Override
173176
public String toString() {
174-
return "Agg[" + getName() + "; " + index.getType() + "]";
177+
return "AGG[" + getName() + "; " + index.getType() + "]";
175178
}
176179

177180
@Override
@@ -359,6 +362,33 @@ public Ordering computeOrderingFromScanComparisons(@Nonnull final ScanComparison
359362
return Ordering.ofOrderingSequence(bindingMapBuilder.build(), orderingSequenceBuilder.build(), isDistinct);
360363
}
361364

365+
@Nullable
366+
@Override
367+
public PullUp.UnificationPullUp prepareForUnification(@Nonnull final PartialMatch partialMatch,
368+
@Nonnull final CorrelationIdentifier topAlias,
369+
@Nonnull final CorrelationIdentifier topCandidateAlias) {
370+
final var regularMatchInfo = partialMatch.getRegularMatchInfo();
371+
if (regularMatchInfo.getRollUpToGroupingValues() != null) {
372+
final var groupingAndAggregateAccessors =
373+
getGroupingAndAggregateAccessors(topCandidateAlias);
374+
final var groupingAccessorValues = groupingAndAggregateAccessors.getLeft();
375+
final var aggregateAccessorValue = groupingAndAggregateAccessors.getRight();
376+
final var allFields =
377+
((Type.Record)selectHavingExpression.getResultValue().getResultType()).getFields();
378+
final var rollUpColumnsBuilder =
379+
ImmutableList.<Column<? extends Value>>builder();
380+
final var numGroupings = regularMatchInfo.getRollUpToGroupingValues().size();
381+
for (int i = 0; i < numGroupings; i++) {
382+
final var field = allFields.get(i);
383+
rollUpColumnsBuilder.add(Column.of(field, groupingAccessorValues.get(i)));
384+
}
385+
rollUpColumnsBuilder.add(Column.of(allFields.get(allFields.size() - 1), aggregateAccessorValue));
386+
return PullUp.forUnification(topAlias, RecordConstructorValue.ofColumns(rollUpColumnsBuilder.build()),
387+
ImmutableSet.of(topCandidateAlias));
388+
}
389+
return null;
390+
}
391+
362392
@Nonnull
363393
@Override
364394
public RecordQueryPlan toEquivalentPlan(@Nonnull final PartialMatch partialMatch,
@@ -369,13 +399,14 @@ public RecordQueryPlan toEquivalentPlan(@Nonnull final PartialMatch partialMatch
369399
final var baseRecordType = Type.Record.fromFieldDescriptorsMap(RecordMetaData.getFieldDescriptorMapFromTypes(recordTypes));
370400

371401
final var selectHavingResultValue = selectHavingExpression.getResultValue();
372-
final var resultType = selectHavingResultValue.getResultType();
402+
final var resultType = (Type.Record)selectHavingResultValue.getResultType();
373403
final var messageDescriptor =
374404
Objects.requireNonNull(TypeRepository.newBuilder()
375405
.addTypeIfNeeded(resultType)
376406
.build()
377407
.getMessageDescriptor(resultType));
378-
final var constraintMaybe = partialMatch.getRegularMatchInfo().getConstraint();
408+
final var regularMatchInfo = partialMatch.getRegularMatchInfo();
409+
final var constraintMaybe = regularMatchInfo.getConstraint();
379410

380411
final var indexEntryConverter = createIndexEntryConverter(messageDescriptor);
381412
final var aggregateIndexScan = new RecordQueryIndexPlan(index.getName(),
@@ -389,12 +420,45 @@ public RecordQueryPlan toEquivalentPlan(@Nonnull final PartialMatch partialMatch
389420
baseRecordType,
390421
QueryPlanConstraint.tautology());
391422

392-
return new RecordQueryAggregateIndexPlan(aggregateIndexScan,
423+
var plan = new RecordQueryAggregateIndexPlan(aggregateIndexScan,
393424
recordTypes.get(0).getName(),
394425
indexEntryConverter,
395426
selectHavingResultValue,
396427
groupByResultValue,
397428
constraintMaybe);
429+
if (regularMatchInfo.getRollUpToGroupingValues() != null) {
430+
//
431+
// We need to perform a roll up.
432+
//
433+
final var aggregateIndexScanReference = memoizer.memoizePlans(plan);
434+
final var aggregateIndexScanAlias = Quantifier.uniqueId();
435+
436+
//final var recordValues = Values.deconstructRecord(recordValue);
437+
final var groupingAndAggregateAccessors =
438+
getGroupingAndAggregateAccessors(regularMatchInfo.getRollUpToGroupingValues().size(),
439+
aggregateIndexScanAlias);
440+
final var groupingAccessorValues = groupingAndAggregateAccessors.getLeft();
441+
final var aggregateAccessorValue = groupingAndAggregateAccessors.getRight();
442+
final var allFields = resultType.getFields();
443+
final var rollUpGroupingColumnsBuilder =
444+
ImmutableList.<Column<? extends Value>>builder();
445+
for (int i = 0; i < groupingAccessorValues.size(); i++) {
446+
final var field = allFields.get(i);
447+
rollUpGroupingColumnsBuilder.add(Column.of(field, groupingAccessorValues.get(i)));
448+
}
449+
450+
final var rollUpAggregateValueOptional =
451+
AggregateIndexExpansionVisitor.rollUpAggregateValueMaybe(index,
452+
aggregateAccessorValue);
453+
454+
final var aggregateIndexScanQuantifier =
455+
Quantifier.physical(aggregateIndexScanReference, aggregateIndexScanAlias);
456+
457+
return RecordQueryStreamingAggregationPlan.ofFlattened(aggregateIndexScanQuantifier,
458+
RecordConstructorValue.ofColumns(rollUpGroupingColumnsBuilder.build(), resultType.isNullable()),
459+
rollUpAggregateValueOptional.orElseThrow(() -> new RecordCoreException("unknown rollup operation")));
460+
}
461+
return plan;
398462
}
399463

400464
@Nonnull
@@ -411,19 +475,22 @@ protected int getGroupingCount() {
411475
}
412476

413477
@Nonnull
414-
public NonnullPair<List<Value>, Value> getGroupingAndAggregateAccessors() {
478+
public NonnullPair<List<Value>, Value> getGroupingAndAggregateAccessors(@Nonnull final CorrelationIdentifier alias) {
479+
return getGroupingAndAggregateAccessors(getGroupingCount(), alias);
480+
}
481+
482+
@Nonnull
483+
public NonnullPair<List<Value>, Value> getGroupingAndAggregateAccessors(final int numGroupings,
484+
@Nonnull final CorrelationIdentifier alias) {
415485
final var selectHavingResultValue = selectHavingExpression.getResultValue();
416486
final var selectHavingResultType = (Type.Record)selectHavingResultValue.getResultType();
417487
final var unbasedResultValue =
418-
QuantifiedObjectValue.of(Quantifier.current(), selectHavingResultType);
419-
420-
final var groupingCount = getGroupingCount();
421-
Verify.verify(selectHavingResultType.getFields().size() >= groupingCount);
422-
488+
QuantifiedObjectValue.of(alias, selectHavingResultType);
423489
final var deconstructedRecord = Values.deconstructRecord(unbasedResultValue);
490+
Verify.verify(deconstructedRecord.size() > numGroupings);
424491
final var groupingAccessorValues =
425-
ImmutableList.copyOf(deconstructedRecord.subList(0, groupingCount));
426-
final var aggregateAccessorValue = deconstructedRecord.get(groupingCount);
492+
ImmutableList.copyOf(deconstructedRecord.subList(0, numGroupings));
493+
final var aggregateAccessorValue = deconstructedRecord.get(deconstructedRecord.size() - 1);
427494
return NonnullPair.of(groupingAccessorValues, aggregateAccessorValue);
428495
}
429496

‎fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/BitmapAggregateIndexExpansionVisitor.java

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,6 @@ protected NonnullPair<Quantifier, List<Placeholder>> constructGroupBy(@Nonnull f
8585
throw new UnsupportedOperationException("unable to plan group by with non-field value " + groupedValue);
8686
}
8787

88-
8988
final var bitmapConstructAggFunc = FunctionCatalog.getFunctionSingleton(NumericAggregationValue.BitmapConstructAggFn.class).orElseThrow();
9089
final var bitmapBitPositionFunc = FunctionCatalog.getFunctionSingleton(ArithmeticValue.BitmapBitPositionFn.class).orElseThrow();
9190
final String sizeArgument = index.getOption(IndexOptions.BITMAP_VALUE_ENTRY_SIZE_OPTION);
@@ -96,7 +95,6 @@ protected NonnullPair<Quantifier, List<Placeholder>> constructGroupBy(@Nonnull f
9695
// add an RCV column representing the grouping columns as the first result set column
9796
// also, make sure to set the field type names correctly for each field value in the grouping keys RCV.
9897

99-
10098
final var groupingValues = baseExpansion.getResultColumns().subList(0, groupingKeyExpression.getGroupingCount())
10199
.stream()
102100
.map(Column::getValue)

‎fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/MatchCandidate.java

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
import com.apple.foundationdb.record.query.plan.cascades.typing.Type;
3737
import com.apple.foundationdb.record.query.plan.cascades.values.Value;
3838
import com.apple.foundationdb.record.query.plan.cascades.values.simplification.OrderingValueComputationRuleSet;
39+
import com.apple.foundationdb.record.query.plan.cascades.values.translation.PullUp;
3940
import com.apple.foundationdb.record.query.plan.plans.RecordQueryPlan;
4041
import com.apple.foundationdb.record.util.pair.NonnullPair;
4142
import com.google.common.base.Verify;
@@ -181,6 +182,13 @@ Ordering computeOrderingFromScanComparisons(@Nonnull ScanComparisons scanCompari
181182
boolean isReverse,
182183
boolean isDistinct);
183184

185+
@Nullable
186+
default PullUp.UnificationPullUp prepareForUnification(@Nonnull final PartialMatch partialMatch,
187+
@Nonnull final CorrelationIdentifier topAlias,
188+
@Nonnull final CorrelationIdentifier topCandidateAlias) {
189+
return null;
190+
}
191+
184192
/**
185193
* Creates a {@link RecordQueryPlan} that represents a scan over the materialized candidate data.
186194
* @param partialMatch the match to be used

0 commit comments

Comments
 (0)