Skip to content

Commit b9cda6b

Browse files
committed
roll ups work
1 parent 813f414 commit b9cda6b

20 files changed

+483
-140
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
@@ -34,6 +34,7 @@
3434
import com.apple.foundationdb.record.query.plan.cascades.expressions.SelectExpression;
3535
import com.apple.foundationdb.record.query.plan.cascades.predicates.Placeholder;
3636
import com.apple.foundationdb.record.query.plan.cascades.predicates.PredicateWithValueAndRanges;
37+
import com.apple.foundationdb.record.query.plan.cascades.values.AggregateValue;
3738
import com.apple.foundationdb.record.query.plan.cascades.values.ArithmeticValue;
3839
import com.apple.foundationdb.record.query.plan.cascades.values.CountValue;
3940
import com.apple.foundationdb.record.query.plan.cascades.values.EmptyValue;
@@ -61,6 +62,7 @@
6162
import java.util.List;
6263
import java.util.Map;
6364
import java.util.Objects;
65+
import java.util.Optional;
6466
import java.util.function.Supplier;
6567
import java.util.stream.Stream;
6668

@@ -72,7 +74,12 @@
7274
public class AggregateIndexExpansionVisitor extends KeyExpressionExpansionVisitor
7375
implements ExpansionVisitor<KeyExpressionExpansionVisitor.VisitorState> {
7476
@Nonnull
75-
static final Supplier<Map<String, BuiltInFunction<? extends Value>>> aggregateMap = Suppliers.memoize(AggregateIndexExpansionVisitor::computeAggregateMap);
77+
static final Supplier<Map<String, BuiltInFunction<? extends Value>>> aggregateMap =
78+
Suppliers.memoize(AggregateIndexExpansionVisitor::computeAggregateMap);
79+
80+
@Nonnull
81+
static final Supplier<Map<String, BuiltInFunction<? extends Value>>> rollUpAggregateMap =
82+
Suppliers.memoize(AggregateIndexExpansionVisitor::computeRollUpAggregateMap);
7683

7784
@Nonnull
7885
protected final Index index;
@@ -92,7 +99,8 @@ public class AggregateIndexExpansionVisitor extends KeyExpressionExpansionVisito
9299
* @param recordTypes The indexed record types.
93100
*/
94101
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()));
102+
Preconditions.checkArgument(IndexTypes.BITMAP_VALUE.equals(index.getType()) ||
103+
aggregateMap.get().containsKey(index.getType()));
96104
Preconditions.checkArgument(index.getRootExpression() instanceof GroupingKeyExpression);
97105
this.index = index;
98106
this.groupingKeyExpression = ((GroupingKeyExpression)index.getRootExpression());
@@ -338,6 +346,12 @@ private ConstructSelectHavingResult constructSelectHaving(@Nonnull final Quantif
338346
finalPlaceholders, currentGroupingValues);
339347
}
340348

349+
@Nonnull
350+
public static Optional<AggregateValue> aggregateValue(@Nonnull final Index index, @Nonnull final Value argument) {
351+
return Optional.of((AggregateValue)aggregateMap.get()
352+
.get(index.getType()).encapsulate(ImmutableList.of(argument)));
353+
}
354+
341355
@Nonnull
342356
private static Map<String, BuiltInFunction<? extends Value>> computeAggregateMap() {
343357
final ImmutableMap.Builder<String, BuiltInFunction<? extends Value>> mapBuilder = ImmutableMap.builder();
@@ -353,6 +367,28 @@ private static Map<String, BuiltInFunction<? extends Value>> computeAggregateMap
353367
return mapBuilder.build();
354368
}
355369

370+
@Nonnull
371+
public static Optional<AggregateValue> rollUpAggregateValueMaybe(@Nonnull final Index index, @Nonnull final Value argument) {
372+
return Optional.ofNullable(rollUpAggregateMap.get()
373+
.get(index.getType()))
374+
.map(fn -> (AggregateValue)fn.encapsulate(ImmutableList.of(argument)));
375+
}
376+
377+
@Nonnull
378+
private static Map<String, BuiltInFunction<? extends Value>> computeRollUpAggregateMap() {
379+
final ImmutableMap.Builder<String, BuiltInFunction<? extends Value>> mapBuilder = ImmutableMap.builder();
380+
mapBuilder.put(IndexTypes.MAX_EVER_LONG, new NumericAggregationValue.MaxFn());
381+
mapBuilder.put(IndexTypes.MIN_EVER_LONG, new NumericAggregationValue.MinFn());
382+
// mapBuilder.put(IndexTypes.MAX_EVER_TUPLE, TODO);
383+
// mapBuilder.put(IndexTypes.MIN_EVER_TUPLE, TODO);
384+
mapBuilder.put(IndexTypes.SUM, new NumericAggregationValue.SumFn());
385+
mapBuilder.put(IndexTypes.COUNT, new NumericAggregationValue.SumFn());
386+
mapBuilder.put(IndexTypes.COUNT_NOT_NULL, new NumericAggregationValue.SumFn());
387+
mapBuilder.put(IndexTypes.PERMUTED_MAX, new NumericAggregationValue.MaxFn());
388+
mapBuilder.put(IndexTypes.PERMUTED_MIN, new NumericAggregationValue.MinFn());
389+
return mapBuilder.build();
390+
}
391+
356392
private static class ConstructSelectHavingResult {
357393
@Nonnull
358394
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
@@ -43,13 +43,16 @@
4343
import com.apple.foundationdb.record.query.plan.cascades.typing.TypeRepository;
4444
import com.apple.foundationdb.record.query.plan.cascades.values.FieldValue;
4545
import com.apple.foundationdb.record.query.plan.cascades.values.QuantifiedObjectValue;
46+
import com.apple.foundationdb.record.query.plan.cascades.values.RecordConstructorValue;
4647
import com.apple.foundationdb.record.query.plan.cascades.values.Value;
4748
import com.apple.foundationdb.record.query.plan.cascades.values.Values;
4849
import com.apple.foundationdb.record.query.plan.cascades.values.simplification.OrderingValueComputationRuleSet;
50+
import com.apple.foundationdb.record.query.plan.cascades.values.translation.PullUp;
4951
import com.apple.foundationdb.record.query.plan.plans.RecordQueryAggregateIndexPlan;
5052
import com.apple.foundationdb.record.query.plan.plans.RecordQueryFetchFromPartialRecordPlan;
5153
import com.apple.foundationdb.record.query.plan.plans.RecordQueryIndexPlan;
5254
import com.apple.foundationdb.record.query.plan.plans.RecordQueryPlan;
55+
import com.apple.foundationdb.record.query.plan.plans.RecordQueryStreamingAggregationPlan;
5356
import com.apple.foundationdb.record.util.pair.NonnullPair;
5457
import com.google.common.base.Preconditions;
5558
import com.google.common.base.Verify;
@@ -172,7 +175,7 @@ public KeyExpression getFullKeyExpression() {
172175

173176
@Override
174177
public String toString() {
175-
return "Agg[" + getName() + "; " + index.getType() + "]";
178+
return "AGG[" + getName() + "; " + index.getType() + "]";
176179
}
177180

178181
@Override
@@ -374,6 +377,33 @@ public Ordering computeOrderingFromScanComparisons(@Nonnull final ScanComparison
374377
return Ordering.ofOrderingSequence(bindingMapBuilder.build(), orderingSequenceBuilder.build(), isDistinct);
375378
}
376379

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+
377407
@Nonnull
378408
@Override
379409
public RecordQueryPlan toEquivalentPlan(@Nonnull final PartialMatch partialMatch,
@@ -384,13 +414,14 @@ public RecordQueryPlan toEquivalentPlan(@Nonnull final PartialMatch partialMatch
384414
final var baseRecordType = Type.Record.fromFieldDescriptorsMap(RecordMetaData.getFieldDescriptorMapFromTypes(recordTypes));
385415

386416
final var selectHavingResultValue = selectHavingExpression.getResultValue();
387-
final var resultType = selectHavingResultValue.getResultType();
417+
final var resultType = (Type.Record)selectHavingResultValue.getResultType();
388418
final var messageDescriptor =
389419
Objects.requireNonNull(TypeRepository.newBuilder()
390420
.addTypeIfNeeded(resultType)
391421
.build()
392422
.getMessageDescriptor(resultType));
393-
final var constraintMaybe = partialMatch.getRegularMatchInfo().getConstraint();
423+
final var regularMatchInfo = partialMatch.getRegularMatchInfo();
424+
final var constraintMaybe = regularMatchInfo.getConstraint();
394425

395426
final var indexEntryConverter = createIndexEntryConverter(messageDescriptor);
396427
final var aggregateIndexScan = new RecordQueryIndexPlan(index.getName(),
@@ -404,12 +435,45 @@ public RecordQueryPlan toEquivalentPlan(@Nonnull final PartialMatch partialMatch
404435
baseRecordType,
405436
QueryPlanConstraint.noConstraint());
406437

407-
return new RecordQueryAggregateIndexPlan(aggregateIndexScan,
438+
var plan = new RecordQueryAggregateIndexPlan(aggregateIndexScan,
408439
recordTypes.get(0).getName(),
409440
indexEntryConverter,
410441
selectHavingResultValue,
411442
groupByResultValue,
412443
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;
413477
}
414478

415479
@Nonnull
@@ -426,19 +490,22 @@ protected int getGroupingCount() {
426490
}
427491

428492
@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) {
430500
final var selectHavingResultValue = selectHavingExpression.getResultValue();
431501
final var selectHavingResultType = (Type.Record)selectHavingResultValue.getResultType();
432502
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);
438504
final var deconstructedRecord = Values.deconstructRecord(unbasedResultValue);
505+
Verify.verify(deconstructedRecord.size() > numGroupings);
439506
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);
442509
return NonnullPair.of(groupingAccessorValues, aggregateAccessorValue);
443510
}
444511

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
@@ -87,7 +87,6 @@ protected NonnullPair<Quantifier, List<Placeholder>> constructGroupBy(@Nonnull f
8787
throw new UnsupportedOperationException("unable to plan group by with non-field value " + groupedValue);
8888
}
8989

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

101-
102100
final var groupingValues = baseExpansion.getResultColumns().subList(0, groupingKeyExpression.getGroupingCount())
103101
.stream()
104102
.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
@@ -37,6 +37,7 @@
3737
import com.apple.foundationdb.record.query.plan.cascades.typing.Type;
3838
import com.apple.foundationdb.record.query.plan.cascades.values.Value;
3939
import com.apple.foundationdb.record.query.plan.cascades.values.simplification.OrderingValueComputationRuleSet;
40+
import com.apple.foundationdb.record.query.plan.cascades.values.translation.PullUp;
4041
import com.apple.foundationdb.record.query.plan.plans.RecordQueryPlan;
4142
import com.apple.foundationdb.record.util.pair.NonnullPair;
4243
import com.google.common.base.Verify;
@@ -182,6 +183,13 @@ Ordering computeOrderingFromScanComparisons(@Nonnull ScanComparisons scanCompari
182183
boolean isReverse,
183184
boolean isDistinct);
184185

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

0 commit comments

Comments
 (0)