diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/AggregationOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/AggregationOperator.java index eae9166af53c..9dcc17ee24f6 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/AggregationOperator.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/AggregationOperator.java @@ -38,6 +38,7 @@ import org.apache.pinot.core.query.request.context.QueryContext; import org.apache.pinot.core.startree.executor.StarTreeAggregationExecutor; import org.apache.pinot.segment.spi.datasource.DataSource; +import org.apache.pinot.segment.spi.index.startree.AggregationFunctionColumnPair; import org.apache.pinot.spi.query.QueryScanCostContext; @@ -47,9 +48,7 @@ public class AggregationOperator extends BaseOperator { private static final String EXPLAIN_NAME = "AGGREGATE"; private final QueryContext _queryContext; - private final AggregationFunction[] _aggregationFunctions; - private final BaseProjectOperator _projectOperator; - private final boolean _useStarTree; + private final AggregationInfo _aggregationInfo; private final int _numTotalDocs; private int _numDocsScanned = 0; @@ -80,9 +79,7 @@ public AggregationOperator(QueryContext queryContext, AggregationInfo aggregatio public AggregationOperator(QueryContext queryContext, AggregationInfo aggregationInfo, int numTotalDocs, @Nullable boolean[] nonScanResolvable, @Nullable DataSource[] dataSources) { _queryContext = queryContext; - _aggregationFunctions = queryContext.getAggregationFunctions(); - _projectOperator = aggregationInfo.getProjectOperator(); - _useStarTree = aggregationInfo.isUseStarTree(); + _aggregationInfo = aggregationInfo; _numTotalDocs = numTotalDocs; _nonScanResolvable = nonScanResolvable; _dataSources = dataSources; @@ -91,27 +88,27 @@ public AggregationOperator(QueryContext queryContext, AggregationInfo aggregatio @Override protected AggregationResultsBlock getNextBlock() { // Perform aggregation on all the transform blocks - AggregationExecutor aggregationExecutor; - if (_useStarTree) { - // StarTreeAggregationExecutor doesn't support non-scan results. - aggregationExecutor = new StarTreeAggregationExecutor(_aggregationFunctions); - } else { - aggregationExecutor = new DefaultAggregationExecutor(_aggregationFunctions, resolveNonScanResults()); - } + AggregationFunction[] aggregationFunctions = _aggregationInfo.getFunctions(); + AggregationFunctionColumnPair[] starTreeFunctionColumnPairs = _aggregationInfo.getStarTreeFunctionColumnPairs(); + AggregationExecutor aggregationExecutor = starTreeFunctionColumnPairs != null + // StarTreeAggregationExecutor doesn't support non-scan results. + ? new StarTreeAggregationExecutor(aggregationFunctions, starTreeFunctionColumnPairs) + : new DefaultAggregationExecutor(aggregationFunctions, resolveNonScanResults()); + BaseProjectOperator projectOperator = _aggregationInfo.getProjectOperator(); ValueBlock valueBlock; - while ((valueBlock = _projectOperator.nextBlock()) != null) { + while ((valueBlock = projectOperator.nextBlock()) != null) { _numDocsScanned += valueBlock.getNumDocs(); QueryScanCostContext scanCost = getScanCostContext(); if (scanCost != null) { scanCost.addDocsScanned(valueBlock.getNumDocs()); scanCost.addEntriesScannedPostFilter( - (long) valueBlock.getNumDocs() * _projectOperator.getNumColumnsProjected()); + (long) valueBlock.getNumDocs() * projectOperator.getNumColumnsProjected()); } aggregationExecutor.aggregate(valueBlock); } // Build intermediate result block based on aggregation result from the executor - return new AggregationResultsBlock(_aggregationFunctions, aggregationExecutor.getResult(), _queryContext); + return new AggregationResultsBlock(aggregationFunctions, aggregationExecutor.getResult(), _queryContext); } /// Returns {@code null} when no function is resolvable without scanning, in which case all functions are computed by @@ -124,10 +121,11 @@ private Object[] resolveNonScanResults() { } Objects.requireNonNull(_dataSources); - Object[] nonScanResults = new Object[_aggregationFunctions.length]; - for (int i = 0; i < _aggregationFunctions.length; i++) { + AggregationFunction[] aggregationFunctions = _aggregationInfo.getFunctions(); + Object[] nonScanResults = new Object[aggregationFunctions.length]; + for (int i = 0; i < aggregationFunctions.length; i++) { if (_nonScanResolvable[i]) { - nonScanResults[i] = AggregationFunctionUtils.getAggregationResult(_aggregationFunctions[i], + nonScanResults[i] = AggregationFunctionUtils.getAggregationResult(aggregationFunctions[i], _dataSources[i], _numTotalDocs, EXPLAIN_NAME); } } @@ -136,13 +134,14 @@ private Object[] resolveNonScanResults() { @Override public List> getChildOperators() { - return List.of(_projectOperator); + return List.of(_aggregationInfo.getProjectOperator()); } @Override public ExecutionStatistics getExecutionStatistics() { - long numEntriesScannedInFilter = _projectOperator.getExecutionStatistics().getNumEntriesScannedInFilter(); - long numEntriesScannedPostFilter = (long) _numDocsScanned * _projectOperator.getNumColumnsProjected(); + BaseProjectOperator projectOperator = _aggregationInfo.getProjectOperator(); + long numEntriesScannedInFilter = projectOperator.getExecutionStatistics().getNumEntriesScannedInFilter(); + long numEntriesScannedPostFilter = (long) _numDocsScanned * projectOperator.getNumColumnsProjected(); return new ExecutionStatistics(_numDocsScanned, numEntriesScannedInFilter, numEntriesScannedPostFilter, _numTotalDocs); } @@ -150,10 +149,11 @@ public ExecutionStatistics getExecutionStatistics() { @Override public String toExplainString() { StringBuilder stringBuilder = new StringBuilder(EXPLAIN_NAME).append("(aggregations:"); - if (_aggregationFunctions.length > 0) { - stringBuilder.append(_aggregationFunctions[0].toExplainString()); - for (int i = 1; i < _aggregationFunctions.length; i++) { - stringBuilder.append(", ").append(_aggregationFunctions[i].toExplainString()); + AggregationFunction[] aggregationFunctions = _aggregationInfo.getFunctions(); + if (aggregationFunctions.length > 0) { + stringBuilder.append(aggregationFunctions[0].toExplainString()); + for (int i = 1; i < aggregationFunctions.length; i++) { + stringBuilder.append(", ").append(aggregationFunctions[i].toExplainString()); } } @@ -168,10 +168,11 @@ protected String getExplainName() { @Override protected void explainAttributes(ExplainAttributeBuilder attributeBuilder) { super.explainAttributes(attributeBuilder); - if (_aggregationFunctions.length == 0) { + AggregationFunction[] aggregationFunctions = _aggregationInfo.getFunctions(); + if (aggregationFunctions.length == 0) { return; } - List aggregations = Arrays.stream(_aggregationFunctions) + List aggregations = Arrays.stream(aggregationFunctions) .map(AggregationFunction::toExplainString) .collect(Collectors.toList()); attributeBuilder.putStringList("aggregations", aggregations); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/FilteredAggregationOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/FilteredAggregationOperator.java index ba433e144743..733b33436927 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/FilteredAggregationOperator.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/FilteredAggregationOperator.java @@ -35,6 +35,7 @@ import org.apache.pinot.core.query.aggregation.function.AggregationFunctionUtils.AggregationInfo; import org.apache.pinot.core.query.request.context.QueryContext; import org.apache.pinot.core.startree.executor.StarTreeAggregationExecutor; +import org.apache.pinot.segment.spi.index.startree.AggregationFunctionColumnPair; import org.apache.pinot.spi.query.QueryScanCostContext; @@ -74,14 +75,12 @@ protected AggregationResultsBlock getNextBlock() { for (AggregationInfo aggregationInfo : _aggregationInfos) { AggregationFunction[] aggregationFunctions = aggregationInfo.getFunctions(); - BaseProjectOperator projectOperator = aggregationInfo.getProjectOperator(); - AggregationExecutor aggregationExecutor; - if (aggregationInfo.isUseStarTree()) { - aggregationExecutor = new StarTreeAggregationExecutor(aggregationFunctions); - } else { - aggregationExecutor = new DefaultAggregationExecutor(aggregationFunctions); - } + AggregationFunctionColumnPair[] starTreeFunctionColumnPairs = aggregationInfo.getStarTreeFunctionColumnPairs(); + AggregationExecutor aggregationExecutor = starTreeFunctionColumnPairs != null + ? new StarTreeAggregationExecutor(aggregationFunctions, starTreeFunctionColumnPairs) + : new DefaultAggregationExecutor(aggregationFunctions); + BaseProjectOperator projectOperator = aggregationInfo.getProjectOperator(); ValueBlock valueBlock; int numDocsScanned = 0; while ((valueBlock = projectOperator.nextBlock()) != null) { @@ -97,8 +96,7 @@ protected AggregationResultsBlock getNextBlock() { QueryScanCostContext scanCost = getScanCostContext(); if (scanCost != null) { scanCost.addDocsScanned(numDocsScanned); - scanCost.addEntriesScannedPostFilter( - (long) numDocsScanned * projectOperator.getNumColumnsProjected()); + scanCost.addEntriesScannedPostFilter((long) numDocsScanned * projectOperator.getNumColumnsProjected()); } _numEntriesScannedInFilter += projectOperator.getExecutionStatistics().getNumEntriesScannedInFilter(); _numEntriesScannedPostFilter += (long) numDocsScanned * projectOperator.getNumColumnsProjected(); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/FilteredGroupByOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/FilteredGroupByOperator.java index f910cd0bda6e..af5b84a87138 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/FilteredGroupByOperator.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/FilteredGroupByOperator.java @@ -49,6 +49,7 @@ import org.apache.pinot.core.query.request.context.QueryContext; import org.apache.pinot.core.startree.executor.StarTreeGroupByExecutor; import org.apache.pinot.core.util.GroupByUtils; +import org.apache.pinot.segment.spi.index.startree.AggregationFunctionColumnPair; import org.apache.pinot.spi.query.QueryScanCostContext; import org.apache.pinot.spi.trace.Tracing; import org.slf4j.Logger; @@ -85,7 +86,7 @@ public FilteredGroupByOperator(QueryContext queryContext, List // NOTE: The indexedTable expects that the data schema will have group by columns before aggregation columns int numGroupByExpressions = _groupByExpressions.length; int numAggregationFunctions = _aggregationFunctions.length; - /// Grouping-set queries append a synthetic $groupingId key column after the union group-by columns. + // Grouping-set queries append a synthetic $groupingId key column after the union group-by columns. int numExtraKeyColumns = queryContext.getNumExtraGroupByKeyColumns(); int numKeyColumns = numGroupByExpressions + numExtraKeyColumns; int numColumns = numKeyColumns + numAggregationFunctions; @@ -101,7 +102,7 @@ public FilteredGroupByOperator(QueryContext queryContext, List projectOperator.getResultColumnContext(groupByExpression).getDataType()); } - /// Synthetic grouping-id discriminator column for GROUP BY GROUPING SETS / ROLLUP / CUBE + // Synthetic grouping-id discriminator column for GROUP BY GROUPING SETS / ROLLUP / CUBE if (numExtraKeyColumns > 0) { columnNames[numGroupByExpressions] = GroupingSets.GROUPING_ID_COLUMN; columnDataTypes[numGroupByExpressions] = DataSchema.ColumnDataType.INT; @@ -141,17 +142,13 @@ protected GroupByResultsBlock getNextBlock() { BaseProjectOperator projectOperator = aggregationInfo.getProjectOperator(); // Perform aggregation group-by on all the blocks - DefaultGroupByExecutor groupByExecutor; + AggregationFunctionColumnPair[] starTreeFunctionColumnPairs = aggregationInfo.getStarTreeFunctionColumnPairs(); + DefaultGroupByExecutor groupByExecutor = starTreeFunctionColumnPairs != null + ? new StarTreeGroupByExecutor(_queryContext, aggregationFunctions, _groupByExpressions, projectOperator, + starTreeFunctionColumnPairs, groupKeyGenerator) + : new DefaultGroupByExecutor(_queryContext, aggregationFunctions, _groupByExpressions, projectOperator, + groupKeyGenerator); - if (aggregationInfo.isUseStarTree()) { - groupByExecutor = - new StarTreeGroupByExecutor(_queryContext, aggregationFunctions, _groupByExpressions, projectOperator, - groupKeyGenerator); - } else { - groupByExecutor = - new DefaultGroupByExecutor(_queryContext, aggregationFunctions, _groupByExpressions, projectOperator, - groupKeyGenerator); - } // The group key generator should be shared across all AggregationFunctions so that agg results can be // aligned. Given that filtered aggregations are stored as an iterable of iterables so that all filtered aggs // with the same filter can share transform blocks, rather than a singular flat iterable in the case where @@ -172,8 +169,7 @@ protected GroupByResultsBlock getNextBlock() { QueryScanCostContext scanCost = getScanCostContext(); if (scanCost != null) { scanCost.addDocsScanned(numDocsScanned); - scanCost.addEntriesScannedPostFilter( - (long) numDocsScanned * projectOperator.getNumColumnsProjected()); + scanCost.addEntriesScannedPostFilter((long) numDocsScanned * projectOperator.getNumColumnsProjected()); } _numEntriesScannedInFilter += projectOperator.getExecutionStatistics().getNumEntriesScannedInFilter(); _numEntriesScannedPostFilter += (long) numDocsScanned * projectOperator.getNumColumnsProjected(); @@ -212,11 +208,11 @@ protected GroupByResultsBlock getNextBlock() { boolean unsafeTrim = _queryContext.isUnsafeTrim(); GroupByResultsBlock resultsBlock; - /// Grouping-set queries use a per-set bucketed segment trim (keyed on the $groupingId discriminator) so - /// that a global top-K cannot starve low-magnitude sets such as the grand total. The broker still applies - /// the final ORDER BY + LIMIT across all sets. + // Grouping-set queries use a per-set bucketed segment trim (keyed on the $groupingId discriminator) so + // that a global top-K cannot starve low-magnitude sets such as the grand total. The broker still applies + // the final ORDER BY + LIMIT across all sets. if (_queryContext.isGroupingSets()) { - /// The $groupingId discriminator is the key column immediately after the union group-by columns. + // The $groupingId discriminator is the key column immediately after the union group-by columns. return GroupByUtils.buildGroupingSetsResultsBlock(_queryContext, _dataSchema, groupKeyGenerator, groupByResultHolders, groupKeyGenerator.getNumKeys(), _groupByExpressions.length, numGroupsLimitReached, numGroupsWarningLimitReached); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/GroupByOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/GroupByOperator.java index 50afa03d62a7..735cc90c0bed 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/GroupByOperator.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/GroupByOperator.java @@ -43,6 +43,7 @@ import org.apache.pinot.core.query.request.context.QueryContext; import org.apache.pinot.core.startree.executor.StarTreeGroupByExecutor; import org.apache.pinot.core.util.GroupByUtils; +import org.apache.pinot.segment.spi.index.startree.AggregationFunctionColumnPair; import org.apache.pinot.spi.query.QueryScanCostContext; import org.apache.pinot.spi.trace.Tracing; import org.slf4j.Logger; @@ -56,10 +57,8 @@ public class GroupByOperator extends BaseOperator { private static final String EXPLAIN_NAME = "GROUP_BY"; private final QueryContext _queryContext; - private final AggregationFunction[] _aggregationFunctions; + private final AggregationInfo _aggregationInfo; private final ExpressionContext[] _groupByExpressions; - private final BaseProjectOperator _projectOperator; - private final boolean _useStarTree; private final long _numTotalDocs; private final DataSchema _dataSchema; @@ -68,17 +67,17 @@ public class GroupByOperator extends BaseOperator { public GroupByOperator(QueryContext queryContext, AggregationInfo aggregationInfo, long numTotalDocs) { assert queryContext.getAggregationFunctions() != null && queryContext.getGroupByExpressions() != null; _queryContext = queryContext; - _aggregationFunctions = queryContext.getAggregationFunctions(); + _aggregationInfo = aggregationInfo; _groupByExpressions = queryContext.getGroupByExpressions().toArray(new ExpressionContext[0]); - _projectOperator = aggregationInfo.getProjectOperator(); - _useStarTree = aggregationInfo.isUseStarTree(); _numTotalDocs = numTotalDocs; // NOTE: The indexedTable expects that the data schema will have group by columns before aggregation columns + AggregationFunction[] aggregationFunctions = aggregationInfo.getFunctions(); + BaseProjectOperator projectOperator = aggregationInfo.getProjectOperator(); int numGroupByExpressions = _groupByExpressions.length; - int numAggregationFunctions = _aggregationFunctions.length; - /// Grouping-set queries append a synthetic $groupingId key column after the union group-by columns (the - /// per-set bitmask discriminator); the key columns thus precede the aggregation columns. + int numAggregationFunctions = aggregationFunctions.length; + // Grouping-set queries append a synthetic $groupingId key column after the union group-by columns (the per-set + // bitmask discriminator); the key columns thus precede the aggregation columns. int numExtraKeyColumns = _queryContext.getNumExtraGroupByKeyColumns(); int numKeyColumns = numGroupByExpressions + numExtraKeyColumns; int numColumns = numKeyColumns + numAggregationFunctions; @@ -90,10 +89,10 @@ public GroupByOperator(QueryContext queryContext, AggregationInfo aggregationInf ExpressionContext groupByExpression = _groupByExpressions[i]; columnNames[i] = groupByExpression.toString(); columnDataTypes[i] = DataSchema.ColumnDataType.fromDataTypeSV( - _projectOperator.getResultColumnContext(groupByExpression).getDataType()); + projectOperator.getResultColumnContext(groupByExpression).getDataType()); } - /// Synthetic grouping-id discriminator column for GROUP BY GROUPING SETS / ROLLUP / CUBE + // Synthetic grouping-id discriminator column for GROUP BY GROUPING SETS / ROLLUP / CUBE if (numExtraKeyColumns > 0) { columnNames[numGroupByExpressions] = GroupingSets.GROUPING_ID_COLUMN; columnDataTypes[numGroupByExpressions] = DataSchema.ColumnDataType.INT; @@ -101,7 +100,7 @@ public GroupByOperator(QueryContext queryContext, AggregationInfo aggregationInf // Extract column names and data types for aggregation functions for (int i = 0; i < numAggregationFunctions; i++) { - AggregationFunction aggregationFunction = _aggregationFunctions[i]; + AggregationFunction aggregationFunction = aggregationFunctions[i]; int index = numKeyColumns + i; columnNames[index] = aggregationFunction.getResultColumnName(); columnDataTypes[index] = aggregationFunction.getIntermediateResultColumnType(); @@ -118,22 +117,21 @@ protected GroupByResultsBlock getNextBlock() { } // Perform aggregation group-by on all the blocks - GroupByExecutor groupByExecutor; // TODO: pass trimGroupSize to executor, who creates the result holder - if (_useStarTree) { - groupByExecutor = new StarTreeGroupByExecutor(_queryContext, _groupByExpressions, _projectOperator); - } else { - groupByExecutor = new DefaultGroupByExecutor(_queryContext, _groupByExpressions, _projectOperator); - } + BaseProjectOperator projectOperator = _aggregationInfo.getProjectOperator(); + AggregationFunctionColumnPair[] starTreeFunctionColumnPairs = _aggregationInfo.getStarTreeFunctionColumnPairs(); + GroupByExecutor groupByExecutor = starTreeFunctionColumnPairs != null + ? new StarTreeGroupByExecutor(_queryContext, _groupByExpressions, projectOperator, starTreeFunctionColumnPairs) + : new DefaultGroupByExecutor(_queryContext, _groupByExpressions, projectOperator); ValueBlock valueBlock; - while ((valueBlock = _projectOperator.nextBlock()) != null) { + while ((valueBlock = projectOperator.nextBlock()) != null) { _numDocsScanned += valueBlock.getNumDocs(); QueryScanCostContext scanCost = getScanCostContext(); if (scanCost != null) { scanCost.addDocsScanned(valueBlock.getNumDocs()); scanCost.addEntriesScannedPostFilter( - (long) valueBlock.getNumDocs() * _projectOperator.getNumColumnsProjected()); + (long) valueBlock.getNumDocs() * projectOperator.getNumColumnsProjected()); } groupByExecutor.process(valueBlock); } @@ -162,11 +160,11 @@ protected GroupByResultsBlock getNextBlock() { boolean unsafeTrim = _queryContext.isUnsafeTrim(); GroupByResultsBlock resultsBlock; - /// Grouping-set queries use a per-set bucketed segment trim (keyed on the $groupingId discriminator) so - /// that a global top-K cannot starve low-magnitude sets such as the grand total. The broker still applies - /// the final ORDER BY + LIMIT across all sets. + // Grouping-set queries use a per-set bucketed segment trim (keyed on the $groupingId discriminator) so + // that a global top-K cannot starve low-magnitude sets such as the grand total. The broker still applies + // the final ORDER BY + LIMIT across all sets. if (_queryContext.isGroupingSets()) { - /// The $groupingId discriminator is the key column immediately after the union group-by columns. + // The $groupingId discriminator is the key column immediately after the union group-by columns. return GroupByUtils.buildGroupingSetsResultsBlock(_queryContext, _dataSchema, groupByExecutor.getGroupKeyGenerator(), groupByExecutor.getGroupByResultHolders(), groupByExecutor.getNumGroups(), _groupByExpressions.length, numGroupsLimitReached, @@ -213,13 +211,14 @@ protected GroupByResultsBlock getNextBlock() { @Override public List getChildOperators() { - return List.of(_projectOperator); + return List.of(_aggregationInfo.getProjectOperator()); } @Override public ExecutionStatistics getExecutionStatistics() { - long numEntriesScannedInFilter = _projectOperator.getExecutionStatistics().getNumEntriesScannedInFilter(); - long numEntriesScannedPostFilter = (long) _numDocsScanned * _projectOperator.getNumColumnsProjected(); + BaseProjectOperator projectOperator = _aggregationInfo.getProjectOperator(); + long numEntriesScannedInFilter = projectOperator.getExecutionStatistics().getNumEntriesScannedInFilter(); + long numEntriesScannedPostFilter = (long) _numDocsScanned * projectOperator.getNumColumnsProjected(); return new ExecutionStatistics(_numDocsScanned, numEntriesScannedInFilter, numEntriesScannedPostFilter, _numTotalDocs); } @@ -235,10 +234,11 @@ public String toExplainString() { } stringBuilder.append(", aggregations:"); - if (_aggregationFunctions.length > 0) { - stringBuilder.append(_aggregationFunctions[0].toExplainString()); - for (int i = 1; i < _aggregationFunctions.length; i++) { - stringBuilder.append(", ").append(_aggregationFunctions[i].toExplainString()); + AggregationFunction[] aggregationFunctions = _aggregationInfo.getFunctions(); + if (aggregationFunctions.length > 0) { + stringBuilder.append(aggregationFunctions[0].toExplainString()); + for (int i = 1; i < aggregationFunctions.length; i++) { + stringBuilder.append(", ").append(aggregationFunctions[i].toExplainString()); } } @@ -258,7 +258,7 @@ protected void explainAttributes(ExplainAttributeBuilder attributeBuilder) { .collect(Collectors.toList()); attributeBuilder.putStringList("groupKeys", groupKeys); - List aggregations = Arrays.stream(_aggregationFunctions) + List aggregations = Arrays.stream(_aggregationInfo.getFunctions()) .map(AggregationFunction::toExplainString) .collect(Collectors.toList()); attributeBuilder.putStringList("aggregations", aggregations); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java index f8c240a3f5c4..ef13e71dceb7 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java @@ -80,20 +80,30 @@ /// The `AggregationFunctionUtils` class provides utility methods for aggregation function. @SuppressWarnings({"rawtypes", "unchecked"}) public class AggregationFunctionUtils { - private AggregationFunctionUtils() { } - /// (For Star-Tree) Creates an [AggregationFunctionColumnPair] in stored type from the - /// [AggregationFunction]. Returns `null` if the [AggregationFunction] cannot be represented as an - /// [AggregationFunctionColumnPair] (e.g. has multiple arguments, argument is not column etc.). + /// (For Star-Tree) Resolves the [AggregationFunctionColumnPair] stored by a regular star-tree, or by a null-aware + /// one when `nullHandlingEnabled` is `true`. Returns `null` if the [AggregationFunction] cannot be represented as a + /// pair (e.g. has multiple arguments, argument is not a column etc.). + /// + /// The two only differ for `COUNT`: a regular star-tree stores a single `count__*` of every row, while a null-aware + /// star-tree stores a `count__column` holding the count of that column's non-null values. @Nullable - public static AggregationFunctionColumnPair getStoredFunctionColumnPair(AggregationFunction aggregationFunction) { + public static AggregationFunctionColumnPair getStoredFunctionColumnPair(AggregationFunction aggregationFunction, + boolean nullHandlingEnabled) { AggregationFunctionType functionType = aggregationFunction.getType(); + List inputExpressions = aggregationFunction.getInputExpressions(); if (functionType == AggregationFunctionType.COUNT) { - return AggregationFunctionColumnPair.COUNT_STAR; + // CountAggregationFunction only reports an input expression when null handling is enabled and the argument is a + // non-star identifier or function, which is exactly when nulls have to be excluded from the count. Everything + // else counts every row and maps to COUNT(*). + if (!nullHandlingEnabled || inputExpressions.size() != 1 + || inputExpressions.get(0).getType() != ExpressionContext.Type.IDENTIFIER) { + return AggregationFunctionColumnPair.COUNT_STAR; + } + return AggregationFunctionColumnPair.countColumn(inputExpressions.get(0).getIdentifier()); } - List inputExpressions = aggregationFunction.getInputExpressions(); if (inputExpressions.size() == 1) { ExpressionContext inputExpression = inputExpressions.get(0); if (inputExpression.getType() == ExpressionContext.Type.IDENTIFIER) { @@ -178,7 +188,13 @@ public static Map getBlockValSetMap(AggregationF /// function pair so that the aggregation result column name is consistent with or without star-tree. public static Map getBlockValSetMap( AggregationFunctionColumnPair aggregationFunctionColumnPair, ValueBlock valueBlock) { - ExpressionContext expression = ExpressionContext.forIdentifier(aggregationFunctionColumnPair.getColumn()); + // A COUNT column holds pre-aggregated counts that have to be summed rather than counted, which + // CountAggregationFunction recognizes from the STAR identifier. This also applies to the per-column COUNT of a + // null-aware star-tree, so the key is STAR there as well. + String column = aggregationFunctionColumnPair.getFunctionType() == AggregationFunctionType.COUNT + ? AggregationFunctionColumnPair.STAR + : aggregationFunctionColumnPair.getColumn(); + ExpressionContext expression = ExpressionContext.forIdentifier(column); BlockValSet blockValSet = valueBlock.getBlockValueSet(aggregationFunctionColumnPair.toColumnName()); return Map.of(expression, blockValSet); } @@ -343,13 +359,21 @@ public static Object getConvertedFinalResult(DataTable dataTable, ColumnDataType public static class AggregationInfo { private final AggregationFunction[] _functions; private final BaseProjectOperator _projectOperator; - private final boolean _useStarTree; + @Nullable + private final AggregationFunctionColumnPair[] _starTreeFunctionColumnPairs; + /// Creates the info for aggregations that do not read a star-tree. + public AggregationInfo(AggregationFunction[] functions, BaseProjectOperator projectOperator) { + this(functions, projectOperator, null); + } + + /// Creates the info for aggregations that read the star-tree the given pairs were resolved against, or for + /// aggregations that do not read one at all when `starTreeFunctionColumnPairs` is `null`. public AggregationInfo(AggregationFunction[] functions, BaseProjectOperator projectOperator, - boolean useStarTree) { + @Nullable AggregationFunctionColumnPair[] starTreeFunctionColumnPairs) { _functions = functions; _projectOperator = projectOperator; - _useStarTree = useStarTree; + _starTreeFunctionColumnPairs = starTreeFunctionColumnPairs; } public AggregationFunction[] getFunctions() { @@ -360,8 +384,13 @@ public BaseProjectOperator getProjectOperator() { return _projectOperator; } - public boolean isUseStarTree() { - return _useStarTree; + /// Returns the function-column pairs projected from the star-tree, or `null` when star-tree is not used. + /// + /// These are resolved against the specific star-tree the query was routed to, so the aggregation executors must + /// use them rather than re-deriving the pairs, which cannot tell a regular star-tree from a null-aware one. + @Nullable + public AggregationFunctionColumnPair[] getStarTreeFunctionColumnPairs() { + return _starTreeFunctionColumnPairs; } } @@ -383,17 +412,18 @@ public static AggregationInfo buildAggregationInfo(SegmentContext segmentContext public static AggregationInfo buildAggregationInfoWithStarTree(SegmentContext segmentContext, QueryContext queryContext, AggregationFunction[] aggregationFunctions, @Nullable FilterContext filter, BaseFilterOperator filterOperator, List> predicateEvaluators) { - /// Star-tree stores pre-aggregated values per group key and cannot expand a row across multiple grouping - /// sets, so it cannot serve GROUP BY GROUPING SETS / ROLLUP / CUBE queries. Fall back to the regular path. + // Star-tree stores pre-aggregated values per group key and cannot expand a row across multiple grouping + // sets, so it cannot serve GROUP BY GROUPING SETS / ROLLUP / CUBE queries. Fall back to the regular path. if (queryContext.isGroupingSets()) { return null; } if (!filterOperator.isResultEmpty()) { - BaseProjectOperator projectOperator = + StarTreeUtils.StarTreeProjectPlan projectPlan = StarTreeUtils.createStarTreeBasedProjectOperator(segmentContext.getIndexSegment(), queryContext, aggregationFunctions, filter, predicateEvaluators); - if (projectOperator != null) { - return new AggregationInfo(aggregationFunctions, projectOperator, true); + if (projectPlan != null) { + return new AggregationInfo(aggregationFunctions, projectPlan.getProjectOperator(), + projectPlan.getFunctionColumnPairs()); } } return null; @@ -408,7 +438,7 @@ public static AggregationInfo buildAggregationInfoWithoutStarTree(SegmentContext BaseProjectOperator projectOperator = new ProjectPlanNode(segmentContext, queryContext, expressionsToTransform, DocIdSetPlanNode.MAX_DOC_PER_CALL, filterOperator).run(); - return new AggregationInfo(aggregationFunctions, projectOperator, false); + return new AggregationInfo(aggregationFunctions, projectOperator); } /// Builds {@link AggregationInfo} for aggregations without using star-tree index, projecting only the columns @@ -423,10 +453,9 @@ public static AggregationInfo buildAggregationInfoWithoutStarTree(SegmentContext BaseProjectOperator projectOperator = new ProjectPlanNode(segmentContext, queryContext, expressionsToTransform, DocIdSetPlanNode.MAX_DOC_PER_CALL, filterOperator).run(); - return new AggregationInfo(allFunctions, projectOperator, false); + return new AggregationInfo(allFunctions, projectOperator); } - /// Builds swim-lanes (list of {@link AggregationInfo}) for filtered aggregations. public static List buildFilteredAggregationInfos(SegmentContext segmentContext, QueryContext queryContext) { @@ -444,7 +473,7 @@ public static List buildFilteredAggregationInfos(SegmentContext BaseProjectOperator projectOperator = new ProjectPlanNode(segmentContext, queryContext, expressions, DocIdSetPlanNode.MAX_DOC_PER_CALL, mainFilterOperator).run(); - return List.of(new AggregationInfo(aggregationFunctions, projectOperator, false)); + return List.of(new AggregationInfo(aggregationFunctions, projectOperator)); } // For each aggregation function, check if the aggregation function is a filtered aggregate. If so, populate the diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DefaultGroupByExecutor.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DefaultGroupByExecutor.java index 5bad8a3fc48b..1c2ba3df39cc 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DefaultGroupByExecutor.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DefaultGroupByExecutor.java @@ -68,11 +68,6 @@ public DefaultGroupByExecutor(QueryContext queryContext, ExpressionContext[] gro this(queryContext, queryContext.getAggregationFunctions(), groupByExpressions, projectOperator, null); } - public DefaultGroupByExecutor(QueryContext queryContext, AggregationFunction[] aggregationFunctions, - ExpressionContext[] groupByExpressions, BaseProjectOperator projectOperator) { - this(queryContext, aggregationFunctions, groupByExpressions, projectOperator, null); - } - public DefaultGroupByExecutor(QueryContext queryContext, AggregationFunction[] aggregationFunctions, ExpressionContext[] groupByExpressions, BaseProjectOperator projectOperator, @Nullable GroupKeyGenerator groupKeyGenerator) { @@ -90,8 +85,8 @@ public DefaultGroupByExecutor(QueryContext queryContext, AggregationFunction[] a // isDictionaryEncoded() flag rather than gating on dictionary nullness alone. hasNoDictionaryGroupByExpression |= !columnContext.isDictionaryEncoded(); } - /// Grouping-set queries expand each row into one group per grouping set, so they always use the - /// multi-value (int[][]) executor path even though the union group-by columns are single-valued. + // Grouping-set queries expand each row into one group per grouping set, so they always use the + // multi-value (int[][]) executor path even though the union group-by columns are single-valued. boolean groupingSets = queryContext.isGroupingSets(); _hasMVGroupByExpression = hasMVGroupByExpression || groupingSets; diff --git a/pinot-core/src/main/java/org/apache/pinot/core/startree/StarTreeUtils.java b/pinot-core/src/main/java/org/apache/pinot/core/startree/StarTreeUtils.java index e9352c5799fb..cf02d0427be0 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/startree/StarTreeUtils.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/startree/StarTreeUtils.java @@ -64,12 +64,24 @@ private StarTreeUtils() { @Nullable public static AggregationFunctionColumnPair[] extractAggregationFunctionPairs( AggregationFunction[] aggregationFunctions) { + return extractAggregationFunctionPairs(aggregationFunctions, false); + } + + /// Extracts the [AggregationFunctionColumnPair]s from the given [AggregationFunction]s, resolving them against + /// either a regular or a null-aware star-tree. Returns `null` if any [AggregationFunction] cannot be represented as + /// an [AggregationFunctionColumnPair]. + /// + /// The only pair that differs between the two is `COUNT`: a regular star-tree stores a single `count__*` of every + /// row, while a null-aware star-tree stores a `count__column` holding the count of that column's non-null values. + @Nullable + public static AggregationFunctionColumnPair[] extractAggregationFunctionPairs( + AggregationFunction[] aggregationFunctions, boolean nullHandlingEnabled) { int numAggregationFunctions = aggregationFunctions.length; AggregationFunctionColumnPair[] aggregationFunctionColumnPairs = new AggregationFunctionColumnPair[numAggregationFunctions]; for (int i = 0; i < numAggregationFunctions; i++) { AggregationFunctionColumnPair aggregationFunctionColumnPair = - AggregationFunctionUtils.getStoredFunctionColumnPair(aggregationFunctions[i]); + AggregationFunctionUtils.getStoredFunctionColumnPair(aggregationFunctions[i], nullHandlingEnabled); if (aggregationFunctionColumnPair != null) { aggregationFunctionColumnPairs[i] = aggregationFunctionColumnPair; } else { @@ -327,6 +339,14 @@ private static PredicateEvaluator getPredicateEvaluator(IndexSegment indexSegmen // Do not use star-tree for the following predicates because: // - REGEXP_LIKE: Need to scan the whole dictionary to gather the matching dictionary ids // - TEXT_MATCH/IS_NULL/IS_NOT_NULL: No way to gather the matching dictionary ids + // TODO: Support IS_NULL / IS_NOT_NULL on a null-aware star-tree. + // Nothing in the index prevents it: a null-aware star-tree stores nulls under a reserved dictionary id one + // past the column's last real id, so IS_NULL matches that id alone and IS_NOT_NULL matches every real id. + // Null rows form their own child node, and StarTreeFilterOperator already skips the star node for a + // predicated dimension, so nulls cannot leak in through it (at the cost of enumerating real children for + // IS_NOT_NULL). The gap is that FilterPlanNode answers both straight from the segment's null vector with a + // BitmapBasedFilterOperator and never builds a predicate evaluator, so there is no getMatchingDictIds() to + // call here. Supporting them needs a star-tree specific evaluator that knows the reserved id. case REGEXP_LIKE: case TEXT_MATCH: case IS_NULL: @@ -356,9 +376,39 @@ static PredicateEvaluator toDictionaryBased(PredicateEvaluator evaluator, Predic dataSource.getDataSourceMetadata().getDataType(), null); } - /// Returns a [BaseProjectOperator] when the filter can be solved with star-tree, or `null` otherwise. + /// The star-tree a query was routed to, together with the [AggregationFunctionColumnPair]s resolved against it. + /// + /// The pairs have to travel with the operator because they depend on which star-tree was picked: a null-aware + /// star-tree resolves `COUNT(column)` to `count__column`, while a regular one resolves it to `count__*`. The + /// aggregation executors must read back the same columns that were projected. + public static class StarTreeProjectPlan { + private final BaseProjectOperator _projectOperator; + private final AggregationFunctionColumnPair[] _functionColumnPairs; + + public StarTreeProjectPlan(BaseProjectOperator projectOperator, + AggregationFunctionColumnPair[] functionColumnPairs) { + _projectOperator = projectOperator; + _functionColumnPairs = functionColumnPairs; + } + + public BaseProjectOperator getProjectOperator() { + return _projectOperator; + } + + public AggregationFunctionColumnPair[] getFunctionColumnPairs() { + return _functionColumnPairs; + } + } + + /// Returns a [StarTreeProjectPlan] when the filter can be solved with star-tree, or `null` otherwise. + /// + /// A star-tree is only consistent with one null-handling mode. A regular star-tree folds nulls into the column's + /// default null value and includes them in the pre-aggregation, matching null-handling-off semantics; a null-aware + /// star-tree keeps nulls apart and excludes them, matching null-handling-on semantics. Queries are therefore routed + /// to a star-tree built in the matching mode, except that a null-handling-on query may still fall back to a regular + /// star-tree when none of the columns it touches actually contains a null value. @Nullable - public static BaseProjectOperator createStarTreeBasedProjectOperator(IndexSegment indexSegment, + public static StarTreeProjectPlan createStarTreeBasedProjectOperator(IndexSegment indexSegment, QueryContext queryContext, AggregationFunction[] aggregationFunctions, @Nullable FilterContext filter, List> predicateEvaluators) { List starTrees = indexSegment.getStarTrees(); @@ -366,12 +416,6 @@ public static BaseProjectOperator createStarTreeBasedProjectOperator(IndexSeg return null; } - AggregationFunctionColumnPair[] aggregationFunctionColumnPairs = - extractAggregationFunctionPairs(aggregationFunctions); - if (aggregationFunctionColumnPairs == null) { - return null; - } - Map> predicateEvaluatorsMap = extractPredicateEvaluatorsMap(indexSegment, filter, predicateEvaluators); if (predicateEvaluatorsMap == null) { @@ -383,84 +427,117 @@ public static BaseProjectOperator createStarTreeBasedProjectOperator(IndexSeg .toArray(new ExpressionContext[0]) : null; if (queryContext.isNullHandlingEnabled()) { - // We can still use the star-tree index if there aren't actually any null values in this segment for all the - // metrics being aggregated, all the dimensions being filtered on / grouped by. - for (int i = 0; i < aggregationFunctionColumnPairs.length; i++) { - AggregationFunctionColumnPair aggregationFunctionColumnPair = aggregationFunctionColumnPairs[i]; - if (aggregationFunctionColumnPair == AggregationFunctionColumnPair.COUNT_STAR) { - // COUNT aggregation function returns a non-empty input expressions list only when null handling is enabled - // and the input operand is a non-star identifier or function. - List inputExpressions = aggregationFunctions[i].getInputExpressions(); - if (!inputExpressions.isEmpty()) { - if (inputExpressions.get(0).getType() == ExpressionContext.Type.IDENTIFIER) { - DataSource dataSource = indexSegment.getDataSource(inputExpressions.get(0).getIdentifier()); - if (dataSource.getNullValueVector() != null && !dataSource.getNullValueVector() - .getNullBitmap() - .isEmpty()) { - return null; - } - } - } - // Null handling is irrelevant for COUNT(*), COUNT(literal), COUNT(nonNullColumn) - continue; - } - - String column = aggregationFunctionColumnPair.getColumn(); - DataSource dataSource = indexSegment.getDataSourceNullable(column); - if (dataSource == null) { - LOGGER.debug("Cannot use star-tree index because aggregation column: '{}' does not exist", column); - return null; - } - if (dataSource.getNullValueVector() != null && !dataSource.getNullValueVector().getNullBitmap().isEmpty()) { - LOGGER.debug("Cannot use star-tree index because aggregation column: '{}' has null values", column); - return null; - } - } - - for (String column : predicateEvaluatorsMap.keySet()) { - DataSource dataSource = indexSegment.getDataSourceNullable(column); - if (dataSource == null) { - LOGGER.debug("Cannot use star-tree index because filter column: '{}' does not exist", column); - return null; - } - if (dataSource.getNullValueVector() != null && !dataSource.getNullValueVector().getNullBitmap().isEmpty()) { - LOGGER.debug("Cannot use star-tree index because filter column: '{}' has null values", column); - return null; - } + // A null-aware star-tree pre-aggregates with exactly the semantics the query asks for + StarTreeProjectPlan plan = createProjectPlan(indexSegment, queryContext, starTrees, true, aggregationFunctions, + groupByExpressions, predicateEvaluatorsMap); + if (plan != null) { + return plan; } + } + return createProjectPlan(indexSegment, queryContext, starTrees, false, aggregationFunctions, groupByExpressions, + predicateEvaluatorsMap); + } - Set groupByColumns = new HashSet<>(); - if (groupByExpressions != null) { - for (ExpressionContext groupByExpression : groupByExpressions) { - groupByExpression.getColumns(groupByColumns); - } - } - for (String column : groupByColumns) { - DataSource dataSource = indexSegment.getDataSourceNullable(column); - if (dataSource == null) { - LOGGER.debug("Cannot use star-tree index because group-by column: '{}' does not exist", column); - return null; - } - if (dataSource.getNullValueVector() != null && !dataSource.getNullValueVector().getNullBitmap().isEmpty()) { - LOGGER.debug("Cannot use star-tree index because group-by column: '{}' has null values", column); - return null; - } - } + /// Returns a [StarTreeProjectPlan] built on the first star-tree that both matches `nullAware` and fits the query, + /// or `null` if there is none. + /// + /// Resolves the function-column pairs against the same mode, because a null-aware star-tree stores `COUNT` per + /// column while a regular one stores a single count of every row, and the executors have to read back whichever + /// was projected. + @Nullable + private static StarTreeProjectPlan createProjectPlan(IndexSegment indexSegment, QueryContext queryContext, + List starTrees, boolean nullAware, AggregationFunction[] aggregationFunctions, + @Nullable ExpressionContext[] groupByExpressions, + Map> predicateEvaluatorsMap) { + // Only `COUNT` resolves differently between the two, and never to `null`, so a query that cannot be represented + // as pairs at all fails here for either kind of star-tree + AggregationFunctionColumnPair[] functionColumnPairs = + extractAggregationFunctionPairs(aggregationFunctions, nullAware); + if (functionColumnPairs == null) { + return null; + } + // A regular star-tree folded nulls into the column's default value and counted them, so it can only answer a + // null-handling-on query when nothing the query touches is actually null + if (!nullAware && queryContext.isNullHandlingEnabled() && !hasNoNullValues(indexSegment, aggregationFunctions, + functionColumnPairs, predicateEvaluatorsMap.keySet(), groupByExpressions)) { + return null; } List> aggregations = new ArrayList<>(aggregationFunctions.length); for (int i = 0; i < aggregationFunctions.length; i++) { - aggregations.add(Pair.of(aggregationFunctions[i], aggregationFunctionColumnPairs[i])); + aggregations.add(Pair.of(aggregationFunctions[i], functionColumnPairs[i])); } for (StarTreeV2 starTreeV2 : starTrees) { - if (isFitForStarTree(starTreeV2.getMetadata(), aggregations, groupByExpressions, - predicateEvaluatorsMap.keySet())) { - return new StarTreeProjectPlanNode(queryContext, starTreeV2, aggregationFunctionColumnPairs, groupByExpressions, - predicateEvaluatorsMap).run(); + StarTreeV2Metadata metadata = starTreeV2.getMetadata(); + if (metadata.isNullHandlingEnabled() != nullAware) { + continue; + } + if (isFitForStarTree(metadata, aggregations, groupByExpressions, predicateEvaluatorsMap.keySet())) { + BaseProjectOperator projectOperator = + new StarTreeProjectPlanNode(queryContext, starTreeV2, functionColumnPairs, groupByExpressions, + predicateEvaluatorsMap).run(); + return new StarTreeProjectPlan(projectOperator, functionColumnPairs); } } return null; } + + /// Returns whether none of the columns the query touches contains a null value in this segment, in which case a + /// regular star-tree produces the same result as a null-aware one and can serve a null-handling-on query. + private static boolean hasNoNullValues(IndexSegment indexSegment, AggregationFunction[] aggregationFunctions, + AggregationFunctionColumnPair[] functionColumnPairs, Set predicateColumns, + @Nullable ExpressionContext[] groupByExpressions) { + for (int i = 0; i < functionColumnPairs.length; i++) { + AggregationFunctionColumnPair functionColumnPair = functionColumnPairs[i]; + if (functionColumnPair == AggregationFunctionColumnPair.COUNT_STAR) { + // COUNT aggregation function returns a non-empty input expressions list only when null handling is enabled + // and the input operand is a non-star identifier or function. Null handling is irrelevant for COUNT(*), + // COUNT(literal) and COUNT(nonNullColumn). + List inputExpressions = aggregationFunctions[i].getInputExpressions(); + if (!inputExpressions.isEmpty() && inputExpressions.get(0).getType() == ExpressionContext.Type.IDENTIFIER + && !hasNoNullValues(indexSegment, inputExpressions.get(0).getIdentifier(), "aggregation")) { + return false; + } + continue; + } + + if (!hasNoNullValues(indexSegment, functionColumnPair.getColumn(), "aggregation")) { + return false; + } + } + + for (String column : predicateColumns) { + if (!hasNoNullValues(indexSegment, column, "filter")) { + return false; + } + } + + Set groupByColumns = new HashSet<>(); + if (groupByExpressions != null) { + for (ExpressionContext groupByExpression : groupByExpressions) { + groupByExpression.getColumns(groupByColumns); + } + } + for (String column : groupByColumns) { + if (!hasNoNullValues(indexSegment, column, "group-by")) { + return false; + } + } + return true; + } + + private static boolean hasNoNullValues(IndexSegment indexSegment, String column, String columnRole) { + DataSource dataSource = indexSegment.getDataSourceNullable(column); + if (dataSource == null) { + LOGGER.debug("Cannot use star-tree index because {} column: '{}' does not exist", columnRole, column); + return false; + } + if (dataSource.getNullValueVector() != null && !dataSource.getNullValueVector().getNullBitmap().isEmpty()) { + LOGGER.debug("Cannot use star-tree index because {} column: '{}' has null values", columnRole, column); + return false; + } + return true; + } } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/startree/executor/StarTreeAggregationExecutor.java b/pinot-core/src/main/java/org/apache/pinot/core/startree/executor/StarTreeAggregationExecutor.java index 0874619f5b5f..db1397dbf448 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/startree/executor/StarTreeAggregationExecutor.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/startree/executor/StarTreeAggregationExecutor.java @@ -30,21 +30,21 @@ /// - The column in function context is function-column pair /// - No transform function in aggregation /// - For `COUNT` aggregation function, we need to aggregate on the pre-aggregated column +@SuppressWarnings({"rawtypes", "unchecked"}) public class StarTreeAggregationExecutor extends DefaultAggregationExecutor { private final AggregationFunctionColumnPair[] _aggregationFunctionColumnPairs; - - public StarTreeAggregationExecutor(AggregationFunction[] aggregationFunctions) { + /// Creates an executor over the pre-aggregated columns the query was routed to. + /// + /// `aggregationFunctionColumnPairs` must be the pairs the star-tree project operator was built with, because they + /// depend on which star-tree was picked: a null-aware star-tree resolves `COUNT(column)` to `count__column` while a + /// regular one resolves it to `count__*`. + public StarTreeAggregationExecutor(AggregationFunction[] aggregationFunctions, + AggregationFunctionColumnPair[] aggregationFunctionColumnPairs) { // StarTreeAggregationExecutor doesn't support pre-aggregated results. // So, we don't need to pass pre-aggregated results to the super class. super(aggregationFunctions); - - int numAggregationFunctions = aggregationFunctions.length; - _aggregationFunctionColumnPairs = new AggregationFunctionColumnPair[numAggregationFunctions]; - for (int i = 0; i < numAggregationFunctions; i++) { - _aggregationFunctionColumnPairs[i] = - AggregationFunctionUtils.getStoredFunctionColumnPair(aggregationFunctions[i]); - } + _aggregationFunctionColumnPairs = aggregationFunctionColumnPairs; } @Override diff --git a/pinot-core/src/main/java/org/apache/pinot/core/startree/executor/StarTreeGroupByExecutor.java b/pinot-core/src/main/java/org/apache/pinot/core/startree/executor/StarTreeGroupByExecutor.java index f72a76a547ac..67ccb5fd5a89 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/startree/executor/StarTreeGroupByExecutor.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/startree/executor/StarTreeGroupByExecutor.java @@ -43,27 +43,21 @@ public class StarTreeGroupByExecutor extends DefaultGroupByExecutor { private final AggregationFunctionColumnPair[] _aggregationFunctionColumnPairs; public StarTreeGroupByExecutor(QueryContext queryContext, ExpressionContext[] groupByExpressions, - BaseProjectOperator projectOperator) { - this(queryContext, queryContext.getAggregationFunctions(), groupByExpressions, projectOperator, null); - } - - public StarTreeGroupByExecutor(QueryContext queryContext, AggregationFunction[] aggregationFunctions, - ExpressionContext[] groupByExpressions, BaseProjectOperator projectOperator) { - this(queryContext, aggregationFunctions, groupByExpressions, projectOperator, null); + BaseProjectOperator projectOperator, AggregationFunctionColumnPair[] aggregationFunctionColumnPairs) { + this(queryContext, queryContext.getAggregationFunctions(), groupByExpressions, projectOperator, + aggregationFunctionColumnPairs, null); } + /// Creates an executor over the pre-aggregated columns the query was routed to. + /// + /// `aggregationFunctionColumnPairs` must be the pairs the star-tree project operator was built with, because they + /// depend on which star-tree was picked: a null-aware star-tree resolves `COUNT(column)` to `count__column` while a + /// regular one resolves it to `count__*`. public StarTreeGroupByExecutor(QueryContext queryContext, AggregationFunction[] aggregationFunctions, ExpressionContext[] groupByExpressions, BaseProjectOperator projectOperator, - @Nullable GroupKeyGenerator groupKeyGenerator) { + AggregationFunctionColumnPair[] aggregationFunctionColumnPairs, @Nullable GroupKeyGenerator groupKeyGenerator) { super(queryContext, aggregationFunctions, groupByExpressions, projectOperator, groupKeyGenerator); - - assert aggregationFunctions != null; - int numAggregationFunctions = aggregationFunctions.length; - _aggregationFunctionColumnPairs = new AggregationFunctionColumnPair[numAggregationFunctions]; - for (int i = 0; i < numAggregationFunctions; i++) { - _aggregationFunctionColumnPairs[i] = - AggregationFunctionUtils.getStoredFunctionColumnPair(aggregationFunctions[i]); - } + _aggregationFunctionColumnPairs = aggregationFunctionColumnPairs; } @Override diff --git a/pinot-core/src/test/java/org/apache/pinot/core/startree/StarTreeUtilsTest.java b/pinot-core/src/test/java/org/apache/pinot/core/startree/StarTreeUtilsTest.java index cd18702ad885..2914561fac4d 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/startree/StarTreeUtilsTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/startree/StarTreeUtilsTest.java @@ -25,7 +25,6 @@ import org.apache.commons.io.FileUtils; import org.apache.pinot.common.request.context.ExpressionContext; import org.apache.pinot.common.request.context.predicate.EqPredicate; -import org.apache.pinot.core.operator.BaseProjectOperator; import org.apache.pinot.core.operator.blocks.ValueBlock; import org.apache.pinot.core.operator.filter.predicate.BaseRawValueBasedPredicateEvaluator; import org.apache.pinot.core.operator.filter.predicate.EqualsPredicateEvaluatorFactory; @@ -226,13 +225,14 @@ public void testStarTreeAcceleratesEqualityFilterOnRawWithDictionary() { FilterPlanNode filterPlanNode = new FilterPlanNode(new SegmentContext(_segment), queryContext); filterPlanNode.run(); - BaseProjectOperator operator = StarTreeUtils.createStarTreeBasedProjectOperator(_segment, queryContext, - queryContext.getAggregationFunctions(), queryContext.getFilter(), - filterPlanNode.getPredicateEvaluators()); - assertNotNull(operator, "Star-tree plan expected for EQ on RAW+dict dimension"); + StarTreeUtils.StarTreeProjectPlan projectPlan = + StarTreeUtils.createStarTreeBasedProjectOperator(_segment, queryContext, + queryContext.getAggregationFunctions(), queryContext.getFilter(), + filterPlanNode.getPredicateEvaluators()); + assertNotNull(projectPlan, "Star-tree plan expected for EQ on RAW+dict dimension"); // Traversal must not throw; before the fix, StarTreeFilterOperator.getMatchingDictIds threw UOE here. - ValueBlock block = operator.nextBlock(); + ValueBlock block = projectPlan.getProjectOperator().nextBlock(); assertNotNull(block, "Star-tree traversal returned no block"); // Star-tree yields one aggregated document per matching path — same behavior as DICTIONARY-encoded columns. diff --git a/pinot-core/src/test/java/org/apache/pinot/queries/NullAwareStarTreeQueriesTest.java b/pinot-core/src/test/java/org/apache/pinot/queries/NullAwareStarTreeQueriesTest.java new file mode 100644 index 000000000000..e866a7ae7dc3 --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/queries/NullAwareStarTreeQueriesTest.java @@ -0,0 +1,194 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.queries; + +import java.io.File; +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import org.apache.commons.io.FileUtils; +import org.apache.pinot.common.response.broker.BrokerResponseNative; +import org.apache.pinot.common.response.broker.ResultTable; +import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; +import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; +import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader; +import org.apache.pinot.segment.local.startree.v2.builder.MultipleTreesBuilder; +import org.apache.pinot.segment.spi.ImmutableSegment; +import org.apache.pinot.segment.spi.IndexSegment; +import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; +import org.apache.pinot.spi.config.table.StarTreeIndexConfig; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.data.readers.GenericRow; +import org.apache.pinot.spi.utils.ReadMode; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; + + +/// Queries answered from a null-aware star-tree, where the dimension itself contains nulls. +/// +/// `NullAwareStarTreeBuilderTest` only inspects what the builder stored. These go through the query path instead, +/// which is where the reserved null dictionary id has to survive being read back. +public class NullAwareStarTreeQueriesTest extends BaseQueriesTest { + private static final File INDEX_DIR = new File(FileUtils.getTempDirectory(), "NullAwareStarTreeQueriesTest"); + private static final String RAW_TABLE_NAME = "testTable"; + private static final String SEGMENT_NAME = "testSegment"; + private static final String DIMENSION = "d"; + private static final String METRIC = "m"; + + private static final Map QUERY_OPTIONS = Map.of("enableNullHandling", "true"); + + /// `_indexSegments` holds two copies and the harness queries two instances, so every aggregate is scaled by this. + private static final int SEGMENT_COPIES = 4; + + /// One record per leaf, so each distinct dimension value gets its own pre-aggregated document. + private static final int MAX_LEAF_RECORDS = 1; + + /// `d` is null for two rows, so the null-aware star-tree stores those under the reserved dictionary id. The metric + /// sums are distinct per group so a group picking up the wrong rows is visible in the answer. + private static final Integer[] DIMENSION_VALUES = {1, 1, 1, 1, null, null, 2, 2, 2, 2}; + private static final int[] METRIC_VALUES = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10}; + private static final int SUM_WHERE_D_IS_1 = 1 + 2 + 3 + 4; + private static final int SUM_WHERE_D_IS_NULL = 5 + 6; + private static final int SUM_WHERE_D_IS_2 = 7 + 8 + 9 + 10; + + private IndexSegment _indexSegment; + private List _indexSegments; + + @Override + protected String getFilter() { + return ""; + } + + @Override + protected IndexSegment getIndexSegment() { + return _indexSegment; + } + + @Override + protected List getIndexSegments() { + return _indexSegments; + } + + @BeforeClass + public void setUp() + throws Exception { + FileUtils.deleteDirectory(INDEX_DIR); + + Schema schema = new Schema.SchemaBuilder().setSchemaName(RAW_TABLE_NAME) + .addSingleValueDimension(DIMENSION, DataType.INT) + .addMetric(METRIC, DataType.INT) + .build(); + TableConfig tableConfig = + new TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME).setNullHandlingEnabled(true).build(); + + List rows = new ArrayList<>(DIMENSION_VALUES.length); + for (int i = 0; i < DIMENSION_VALUES.length; i++) { + GenericRow row = new GenericRow(); + row.putValue(DIMENSION, DIMENSION_VALUES[i]); + row.putValue(METRIC, METRIC_VALUES[i]); + rows.add(row); + } + + SegmentGeneratorConfig segmentGeneratorConfig = new SegmentGeneratorConfig(tableConfig, schema); + segmentGeneratorConfig.setTableName(RAW_TABLE_NAME); + segmentGeneratorConfig.setSegmentName(SEGMENT_NAME); + segmentGeneratorConfig.setDefaultNullHandlingEnabled(true); + segmentGeneratorConfig.setOutDir(INDEX_DIR.getPath()); + SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl(); + driver.init(segmentGeneratorConfig, new GenericRowRecordReader(rows)); + driver.build(); + + File indexDir = new File(INDEX_DIR, SEGMENT_NAME); + StarTreeIndexConfig starTreeIndexConfig = + new StarTreeIndexConfig(List.of(DIMENSION), null, List.of("SUM__" + METRIC), null, MAX_LEAF_RECORDS, true); + try (MultipleTreesBuilder builder = new MultipleTreesBuilder(List.of(starTreeIndexConfig), false, indexDir, + MultipleTreesBuilder.BuildMode.OFF_HEAP)) { + builder.build(); + } + + ImmutableSegment segment = ImmutableSegmentLoader.load(indexDir, ReadMode.mmap); + _indexSegment = segment; + _indexSegments = List.of(segment, segment); + } + + @AfterClass + public void tearDown() + throws IOException { + _indexSegment.destroy(); + FileUtils.deleteDirectory(INDEX_DIR); + } + + /// The reserved null dictionary id has no entry in the dictionary the star-tree shares with the segment, so + /// resolving it reads past the end of the dictionary's value buffer. + @Test + public void groupingByANullDimensionReturnsTheNullGroup() { + BrokerResponseNative response = + getBrokerResponse("SELECT " + DIMENSION + ", SUM(" + METRIC + ") FROM testTable GROUP BY " + DIMENSION, + QUERY_OPTIONS); + + assertStarTreeWasUsed(response); + ResultTable resultTable = response.getResultTable(); + assertEquals(sumForGroup(resultTable, 1), (double) SUM_WHERE_D_IS_1 * SEGMENT_COPIES); + assertEquals(sumForGroup(resultTable, 2), (double) SUM_WHERE_D_IS_2 * SEGMENT_COPIES); + assertEquals(sumForGroup(resultTable, null), (double) SUM_WHERE_D_IS_NULL * SEGMENT_COPIES, + "Rows whose dimension is null must form their own group rather than joining a real value's group"); + } + + /// A predicate that is always true over real values is not always true over nulls: it is UNKNOWN there, so the row + /// is not selected. Dropping the predicate as always true loses that, and the star-tree then counts the null rows. + @Test + public void anAlwaysTruePredicateStillExcludesNullRows() { + BrokerResponseNative response = + getBrokerResponse("SELECT SUM(" + METRIC + ") FROM testTable WHERE " + DIMENSION + " <> 99999", QUERY_OPTIONS); + + assertStarTreeWasUsed(response); + assertEquals(response.getResultTable().getRows().get(0)[0], + (double) (SUM_WHERE_D_IS_1 + SUM_WHERE_D_IS_2) * SEGMENT_COPIES, + "A null dimension makes the predicate UNKNOWN, so the row must not be aggregated"); + } + + /// Guards against the checks above passing because the query silently fell back to a raw scan, which would answer + /// correctly and prove nothing. A star-tree reads one pre-aggregated document per group rather than every row. + private static void assertStarTreeWasUsed(BrokerResponseNative response) { + long numRowsScanned = (long) DIMENSION_VALUES.length * SEGMENT_COPIES; + assertTrue(response.getNumDocsScanned() < numRowsScanned, + "Expected the star-tree to be used, but " + response.getNumDocsScanned() + " documents were scanned out of " + + numRowsScanned); + } + + /// Returns the aggregate of the row whose group key matches, or fails when no such group exists. + private static double sumForGroup(ResultTable resultTable, Integer groupKey) { + for (Object[] row : resultTable.getRows()) { + if (groupKey == null ? row[0] == null : groupKey.equals(row[0])) { + return ((Number) row[1]).doubleValue(); + } + } + throw new AssertionError("No group for dimension value: " + groupKey + " in " + resultTable.getRows().size() + + " rows"); + } +} diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/CountValueAggregator.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/CountValueAggregator.java index 2b7e01a65754..498eaf669f00 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/CountValueAggregator.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/CountValueAggregator.java @@ -41,6 +41,17 @@ public Long getInitialAggregatedValue(@Nullable Object rawValue) { return rawValue != null ? 1L : 0L; } + /// The only aggregator that still answers a group with no non-null input itself. + /// + /// `COUNT` is read back from the pre-aggregated column by summing it rather than through the null vector, which + /// [org.apache.pinot.core.query.aggregation.function.CountAggregationFunction] recognizes from the `__STAR__` + /// identifier. Returning `0` is also exactly the placeholder a null would leave behind, so recording the group in + /// the null vector would cost a bit per group and buy nothing. + @Override + public Long getAllNullAggregatedValue() { + return 0L; + } + @Override public Long applyRawValue(Long value, Object rawValue) { return value + 1; diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/DistinctCountBitmapValueAggregator.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/DistinctCountBitmapValueAggregator.java index f1d1d4c161c1..74fbb95f2d85 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/DistinctCountBitmapValueAggregator.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/DistinctCountBitmapValueAggregator.java @@ -41,7 +41,9 @@ public DataType getAggregatedValueType() { @Override public RoaringBitmap getInitialAggregatedValue(Object rawValue) { - // NOTE: rawValue cannot be null because this aggregator can only be used for star-tree index. + // NOTE: rawValue cannot be null because this aggregator can only be used for star-tree index, and the builder + // never passes a null raw value: a null-aware star-tree leaves the aggregated value null until the group sees + // its first non-null input. assert rawValue != null; RoaringBitmap initialValue; if (rawValue instanceof byte[]) { diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/DistinctCountThetaSketchValueAggregator.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/DistinctCountThetaSketchValueAggregator.java index c735ea4f9067..b2e72d99d6f9 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/DistinctCountThetaSketchValueAggregator.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/DistinctCountThetaSketchValueAggregator.java @@ -116,7 +116,9 @@ private void multiItemUpdate(ThetaUnion thetaUnion, Object[] rawValues) { @Override public Object getInitialAggregatedValue(Object rawValue) { - // NOTE: rawValue cannot be null because this aggregator can only be used for star-tree index. + // NOTE: rawValue cannot be null because this aggregator can only be used for star-tree index, and the builder + // never passes a null raw value: a null-aware star-tree leaves the aggregated value null until the group + // sees its first non-null input. assert rawValue != null; ThetaUnion thetaUnion = _setOperationBuilder.buildUnion(); if (rawValue instanceof byte[]) { // Serialized ThetaSketch diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/ValueAggregator.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/ValueAggregator.java index 506f02fc0f79..6664bbf7bb05 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/ValueAggregator.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/ValueAggregator.java @@ -41,6 +41,28 @@ public interface ValueAggregator { /// specified in the schema. A getInitialAggregatedValue(@Nullable R rawValue); + /// Returns the aggregated value of a group whose input values are all null, or `null` to have the star-tree record + /// the group in its null vector instead. + /// + /// Only consulted by null-aware star-trees, which exclude null input values from the pre-aggregation and can + /// therefore produce a group with no values at all. + /// + /// Returning `null` is safe whenever the aggregation function skips null rows while reading the pre-aggregated + /// column, which every aggregation function does apart from `COUNT`. A group recorded in the null vector is never + /// read back, so the placeholder left in the forward index is never deserialized. + /// + /// `COUNT` is the exception and overrides this: it is read back by summing the pre-aggregated column rather than + /// through the null vector, so it answers `0` itself. Every other aggregator takes the default, which keeps an + /// all-null group down to a placeholder plus one null-vector bit instead of a serialized empty sketch. + /// + /// An aggregator whose [#getAggregatedValueType] is `BYTES` must also make [#getMaxAggregatedValueByteSize] account + /// for the value returned here, because the star-tree sizes the variable-length forward index from that and would + /// otherwise under-allocate for a metric whose every group is null. + @Nullable + default A getAllNullAggregatedValue() { + return null; + } + /// Applies a raw value to the current aggregated value. /// /// NOTE: if value is mutable, will directly modify the value. diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/store/StarTreeIndexReader.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/store/StarTreeIndexReader.java index 4e00fcde59f2..e53f33435ed2 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/store/StarTreeIndexReader.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/store/StarTreeIndexReader.java @@ -102,11 +102,15 @@ private void mapBufferEntries(int starTreeId, new StarTreeIndexEntry(indexMap.get(StarTreeIndexMapUtils.STAR_TREE_INDEX_KEY), _dataBuffer, ByteOrder.LITTLE_ENDIAN)); StarTreeV2Metadata starTreeMetadata = _starTreeMetadataList.get(starTreeId); + boolean nullHandlingEnabled = starTreeMetadata.isNullHandlingEnabled(); // Load dimension forward indexes for (String dimension : starTreeMetadata.getDimensionsSplitOrder()) { columnEntries.put(new IndexKey(dimension, StandardIndexes.forward()), new StarTreeIndexEntry( indexMap.get(new StarTreeIndexMapUtils.IndexKey(StarTreeIndexMapUtils.IndexType.FORWARD_INDEX, dimension)), _dataBuffer, ByteOrder.BIG_ENDIAN)); + if (nullHandlingEnabled) { + mapNullValueVectorEntry(columnEntries, indexMap, dimension); + } } // Load metric (function-column pair) forward indexes for (AggregationFunctionColumnPair functionColumnPair : starTreeMetadata.getFunctionColumnPairs()) { @@ -114,6 +118,23 @@ private void mapBufferEntries(int starTreeId, columnEntries.put(new IndexKey(metric, StandardIndexes.forward()), new StarTreeIndexEntry( indexMap.get(new StarTreeIndexMapUtils.IndexKey(StarTreeIndexMapUtils.IndexType.FORWARD_INDEX, metric)), _dataBuffer, ByteOrder.BIG_ENDIAN)); + if (nullHandlingEnabled) { + mapNullValueVectorEntry(columnEntries, indexMap, metric); + } + } + } + + /// Maps the null value vector of a null-aware star-tree column, if it has one. + /// + /// The builder only writes a null value vector for columns that actually contain null values, so a missing entry + /// simply means that the column has none. + private void mapNullValueVectorEntry(Map columnEntries, + Map indexMap, String column) { + StarTreeIndexMapUtils.IndexValue indexValue = + indexMap.get(new StarTreeIndexMapUtils.IndexKey(StarTreeIndexMapUtils.IndexType.NULL_VALUE_VECTOR, column)); + if (indexValue != null) { + columnEntries.put(new IndexKey(column, StandardIndexes.nullValueVector()), + new StarTreeIndexEntry(indexValue, _dataBuffer, ByteOrder.BIG_ENDIAN)); } } diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/StarTreeBuilderUtils.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/StarTreeBuilderUtils.java index 72e73b4cfd8c..fe78850ea8ea 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/StarTreeBuilderUtils.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/StarTreeBuilderUtils.java @@ -263,6 +263,10 @@ public static boolean shouldModifyExistingStarTrees(List newSpecs = builderConfig.getAggregationSpecs(); TreeMap existingSpecs = metadata.getAggregationSpecs(); if (newSpecs.size() != existingSpecs.size()) { diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/builder/BaseSingleTreeBuilder.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/builder/BaseSingleTreeBuilder.java index 456e8f32e5fa..b45f759ce857 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/builder/BaseSingleTreeBuilder.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/builder/BaseSingleTreeBuilder.java @@ -31,12 +31,14 @@ import java.util.TreeMap; import javax.annotation.Nullable; import org.apache.commons.configuration2.Configuration; +import org.apache.commons.lang3.ArrayUtils; import org.apache.pinot.common.request.context.ExpressionContext; import org.apache.pinot.segment.local.aggregator.ValueAggregator; import org.apache.pinot.segment.local.aggregator.ValueAggregatorFactory; import org.apache.pinot.segment.local.segment.creator.impl.fwd.SingleValueFixedByteRawIndexCreator; import org.apache.pinot.segment.local.segment.creator.impl.fwd.SingleValueUnsortedForwardIndexCreator; import org.apache.pinot.segment.local.segment.creator.impl.fwd.SingleValueVarByteRawIndexCreator; +import org.apache.pinot.segment.local.segment.creator.impl.nullvalue.NullValueVectorCreator; import org.apache.pinot.segment.local.segment.readers.PinotSegmentColumnReader; import org.apache.pinot.segment.local.startree.StarTreeBuilderUtils; import org.apache.pinot.segment.local.startree.StarTreeBuilderUtils.TreeNode; @@ -70,6 +72,11 @@ abstract class BaseSingleTreeBuilder implements SingleTreeBuilder { final String[] _dimensionsSplitOrder; final Set _skipStarNodeCreationForDimensions; final PinotSegmentColumnReader[] _dimensionReaders; + // Cardinality of each dimension's segment dictionary. In a null-aware star-tree this doubles as the dictionary id + // reserved for null values, which is why the dimension forward index is then created with cardinality + 1 values. + // Reserving the id one past the last real one keeps nulls sorted after every real value, so they form their own + // tree node instead of folding into the column's default null value. + final int[] _dimensionCardinalities; final int _numMetrics; // Name of the function-column pairs @@ -80,6 +87,7 @@ abstract class BaseSingleTreeBuilder implements SingleTreeBuilder { final AggregationSpec[] _aggregationSpecs; final int _maxLeafRecords; + final boolean _nullHandlingEnabled; final TreeNode _rootNode = getNewNode(); @@ -109,12 +117,14 @@ static class Record { _outputDir = outputDir; _segment = segment; _metadataProperties = metadataProperties; + _nullHandlingEnabled = builderConfig.isNullHandlingEnabled(); List dimensionsSplitOrder = builderConfig.getDimensionsSplitOrder(); _numDimensions = dimensionsSplitOrder.size(); _dimensionsSplitOrder = new String[_numDimensions]; _skipStarNodeCreationForDimensions = new HashSet<>(); _dimensionReaders = new PinotSegmentColumnReader[_numDimensions]; + _dimensionCardinalities = new int[_numDimensions]; Set skipStarNodeCreationForDimensions = builderConfig.getSkipStarNodeCreationForDimensions(); for (int i = 0; i < _numDimensions; i++) { String dimension = dimensionsSplitOrder.get(i); @@ -125,6 +135,7 @@ static class Record { _dimensionReaders[i] = new PinotSegmentColumnReader(segment, dimension); Preconditions.checkState(_dimensionReaders[i].hasDictionary(), "Dimension: " + dimension + " does not have dictionary"); + _dimensionCardinalities[i] = segment.getDictionary(dimension).length(); } TreeMap aggregationSpecs = builderConfig.getAggregationSpecs(); @@ -144,9 +155,11 @@ static class Record { _valueAggregators[index] = ValueAggregatorFactory.getValueAggregator(functionColumnPair.getFunctionType(), arguments); _aggregationSpecs[index] = aggregationSpec; - // Ignore the column for COUNT aggregation function - if (_valueAggregators[index].getAggregationType() != AggregationFunctionType.COUNT) { - String column = functionColumnPair.getColumn(); + // COUNT(*) counts rows rather than values and needs no reader. A null-aware star-tree can additionally store + // COUNT(column), which counts the non-null values of that column and therefore does need one. + String column = functionColumnPair.getColumn(); + if (_valueAggregators[index].getAggregationType() != AggregationFunctionType.COUNT || !column.equals( + AggregationFunctionColumnPair.STAR)) { _metricReaders[index] = new PinotSegmentColumnReader(segment, column); } @@ -209,8 +222,17 @@ abstract Iterator generateRecordsForStarNode(int startDocId, int endDocI /// @return Dimensions (dictionary Ids) for a segment record int[] getSegmentRecordDimensions(int docId) { int[] dimensions = new int[_numDimensions]; - for (int i = 0; i < _numDimensions; i++) { - dimensions[i] = _dimensionReaders[i].getDictId(docId); + if (_nullHandlingEnabled) { + for (int i = 0; i < _numDimensions; i++) { + PinotSegmentColumnReader dimensionReader = _dimensionReaders[i]; + // A null value is stored under the reserved dictionary id instead of the dictionary id of the column's + // default null value, so that null rows are not grouped together with rows holding that default value + dimensions[i] = dimensionReader.isNull(docId) ? _dimensionCardinalities[i] : dimensionReader.getDictId(docId); + } + } else { + for (int i = 0; i < _numDimensions; i++) { + dimensions[i] = _dimensionReaders[i].getDictId(docId); + } } return dimensions; } @@ -223,9 +245,12 @@ Record getSegmentRecord(int docId) { int[] dimensions = getSegmentRecordDimensions(docId); Object[] metrics = new Object[_numMetrics]; for (int i = 0; i < _numMetrics; i++) { - // Ignore the column for COUNT aggregation function - if (_metricReaders[i] != null) { - metrics[i] = _metricReaders[i].getValue(docId); + // Ignore the column for COUNT(*), which has no reader + PinotSegmentColumnReader metricReader = _metricReaders[i]; + if (metricReader != null) { + // A null-aware star-tree excludes null values from the aggregation, so it passes them down as null rather + // than as the column's default null value + metrics[i] = _nullHandlingEnabled && metricReader.isNull(docId) ? null : metricReader.getValue(docId); } } return new Record(dimensions, metrics); @@ -243,24 +268,33 @@ Record mergeSegmentRecord(@Nullable Record aggregatedRecord, Record segmentRecor int[] dimensions = Arrays.copyOf(segmentRecord._dimensions, _numDimensions); Object[] metrics = new Object[_numMetrics]; for (int i = 0; i < _numMetrics; i++) { - Object rawValue = segmentRecord._metrics[i]; - if (rawValue != null) { - metrics[i] = _valueAggregators[i].getInitialAggregatedValue(rawValue); - } else { + if (_metricReaders[i] == null) { + // COUNT(*) has no reader and counts every row assert _valueAggregators[i].getAggregationType() == AggregationFunctionType.COUNT; metrics[i] = 1L; + continue; } + // A null raw value only occurs in a null-aware star-tree, where it is excluded from the aggregation. The + // aggregated value stays null until the group sees its first non-null value. + Object rawValue = segmentRecord._metrics[i]; + metrics[i] = rawValue != null ? _valueAggregators[i].getInitialAggregatedValue(rawValue) : null; } return new Record(dimensions, metrics); } else { for (int i = 0; i < _numMetrics; i++) { - Object rawValue = segmentRecord._metrics[i]; - if (rawValue != null) { - aggregatedRecord._metrics[i] = _valueAggregators[i].applyRawValue(aggregatedRecord._metrics[i], rawValue); - } else { + if (_metricReaders[i] == null) { assert _valueAggregators[i].getAggregationType() == AggregationFunctionType.COUNT; aggregatedRecord._metrics[i] = ((long) aggregatedRecord._metrics[i]) + 1; + continue; + } + Object rawValue = segmentRecord._metrics[i]; + if (rawValue == null) { + continue; } + Object aggregatedValue = aggregatedRecord._metrics[i]; + aggregatedRecord._metrics[i] = + aggregatedValue != null ? _valueAggregators[i].applyRawValue(aggregatedValue, rawValue) + : _valueAggregators[i].getInitialAggregatedValue(rawValue); } return aggregatedRecord; } @@ -278,13 +312,21 @@ Record mergeStarTreeRecord(@Nullable Record aggregatedRecord, Record starTreeRec int[] dimensions = Arrays.copyOf(starTreeRecord._dimensions, _numDimensions); Object[] metrics = new Object[_numMetrics]; for (int i = 0; i < _numMetrics; i++) { - metrics[i] = _valueAggregators[i].cloneAggregatedValue(starTreeRecord._metrics[i]); + // A null value means the group aggregated over no non-null input, which only occurs in a null-aware star-tree + Object value = starTreeRecord._metrics[i]; + metrics[i] = value != null ? _valueAggregators[i].cloneAggregatedValue(value) : null; } return new Record(dimensions, metrics); } else { for (int i = 0; i < _numMetrics; i++) { + Object value = starTreeRecord._metrics[i]; + if (value == null) { + continue; + } + Object aggregatedValue = aggregatedRecord._metrics[i]; aggregatedRecord._metrics[i] = - _valueAggregators[i].applyAggregatedValue(aggregatedRecord._metrics[i], starTreeRecord._metrics[i]); + aggregatedValue != null ? _valueAggregators[i].applyAggregatedValue(aggregatedValue, value) + : _valueAggregators[i].cloneAggregatedValue(value); } return aggregatedRecord; } @@ -455,10 +497,27 @@ private void createForwardIndexes() SingleValueUnsortedForwardIndexCreator[] dimensionIndexCreators = new SingleValueUnsortedForwardIndexCreator[_numDimensions]; for (int i = 0; i < _numDimensions; i++) { - String dimension = _dimensionsSplitOrder[i]; - int cardinality = _segment.getDictionary(dimension).length(); + // A null-aware star-tree reserves one extra dictionary id per dimension for null values, which may widen the + // forward index by one bit per value. StarTreeLoaderUtils applies the same adjustment when reading it back. + int numValues = _nullHandlingEnabled ? _dimensionCardinalities[i] + 1 : _dimensionCardinalities[i]; dimensionIndexCreators[i] = - new SingleValueUnsortedForwardIndexCreator(_outputDir, _dimensionsSplitOrder[i], cardinality, _numDocs); + new SingleValueUnsortedForwardIndexCreator(_outputDir, _dimensionsSplitOrder[i], numValues, _numDocs); + } + + // Null vectors are only created for a null-aware star-tree. Dimensions need one as well as metrics: the reserved + // null dictionary id is out of range for the segment dictionary the star-tree shares, so the query side relies on + // the null vector rather than on the stored dictionary id. + NullValueVectorCreator[] dimensionNullValueVectorCreators = + _nullHandlingEnabled ? new NullValueVectorCreator[_numDimensions] : null; + NullValueVectorCreator[] metricNullValueVectorCreators = + _nullHandlingEnabled ? new NullValueVectorCreator[_numMetrics] : null; + if (_nullHandlingEnabled) { + for (int i = 0; i < _numDimensions; i++) { + dimensionNullValueVectorCreators[i] = new NullValueVectorCreator(_outputDir, _dimensionsSplitOrder[i]); + } + for (int i = 0; i < _numMetrics; i++) { + metricNullValueVectorCreators[i] = new NullValueVectorCreator(_outputDir, _metrics[i]); + } } ForwardIndexCreator[] metricIndexCreators = new ForwardIndexCreator[_numMetrics]; @@ -486,30 +545,36 @@ private void createForwardIndexes() for (int docId = 0; docId < _numDocs; docId++) { Record record = getStarTreeRecord(docId); for (int i = 0; i < _numDimensions; i++) { - dimensionIndexCreators[i].putDictId(record._dimensions[i]); + int dictId = record._dimensions[i]; + dimensionIndexCreators[i].putDictId(dictId); + // Star records store STAR_IN_FORWARD_INDEX (0) for the starred dimension, which never collides with the + // reserved null dictionary id, so star records are correctly left out of the null vector + if (dimensionNullValueVectorCreators != null && dictId == _dimensionCardinalities[i]) { + dimensionNullValueVectorCreators[i].setNull(docId); + } } for (int i = 0; i < _numMetrics; i++) { ValueAggregator valueAggregator = _valueAggregators[i]; - ForwardIndexCreator metricIndexCreator = metricIndexCreators[i]; - switch (valueAggregator.getAggregatedValueType()) { - case INT: - metricIndexCreator.putInt((int) record._metrics[i]); - break; - case LONG: - metricIndexCreator.putLong((long) record._metrics[i]); - break; - case FLOAT: - metricIndexCreator.putFloat((float) record._metrics[i]); - break; - case DOUBLE: - metricIndexCreator.putDouble((double) record._metrics[i]); - break; - case BYTES: - metricIndexCreator.putBytes(valueAggregator.serializeAggregatedValue(record._metrics[i])); - break; - default: - throw new IllegalStateException(); + Object value = record._metrics[i]; + if (value == null) { + // The group aggregated over no non-null input. Only COUNT still has a well-defined result of its own + // (0); every other aggregator answers SQL NULL and gets a placeholder in the forward index masked by + // the null vector. + assert _nullHandlingEnabled; + value = valueAggregator.getAllNullAggregatedValue(); + if (value == null) { + metricNullValueVectorCreators[i].setNull(docId); + } } + putMetricValue(metricIndexCreators[i], valueAggregator, value); + } + } + if (_nullHandlingEnabled) { + for (NullValueVectorCreator nullValueVectorCreator : dimensionNullValueVectorCreators) { + nullValueVectorCreator.seal(); + } + for (NullValueVectorCreator nullValueVectorCreator : metricNullValueVectorCreators) { + nullValueVectorCreator.seal(); } } } catch (Exception e) { @@ -543,6 +608,35 @@ private void createForwardIndexes() } } + /// Writes an aggregated metric value into the forward index. + /// + /// A `null` value is a group that aggregates to SQL `NULL`; it is recorded in the metric's null vector and stored + /// here as the aggregated type's zero value, which the query side never reads. + private static void putMetricValue(ForwardIndexCreator metricIndexCreator, ValueAggregator valueAggregator, + @Nullable Object value) + throws IOException { + switch (valueAggregator.getAggregatedValueType()) { + case INT: + metricIndexCreator.putInt(value != null ? (int) value : 0); + break; + case LONG: + metricIndexCreator.putLong(value != null ? (long) value : 0L); + break; + case FLOAT: + metricIndexCreator.putFloat(value != null ? (float) value : 0f); + break; + case DOUBLE: + metricIndexCreator.putDouble(value != null ? (double) value : 0d); + break; + case BYTES: + metricIndexCreator.putBytes( + value != null ? valueAggregator.serializeAggregatedValue(value) : ArrayUtils.EMPTY_BYTE_ARRAY); + break; + default: + throw new IllegalStateException(); + } + } + @Override public void close() throws IOException { diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/builder/OffHeapSingleTreeBuilder.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/builder/OffHeapSingleTreeBuilder.java index f6ab03fa5805..8ae003075ea2 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/builder/OffHeapSingleTreeBuilder.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/builder/OffHeapSingleTreeBuilder.java @@ -35,6 +35,7 @@ import org.apache.pinot.segment.spi.ImmutableSegment; import org.apache.pinot.segment.spi.index.startree.StarTreeV2Constants; import org.apache.pinot.segment.spi.memory.PinotDataBuffer; +import org.roaringbitmap.RoaringBitmap; /// The `OffHeapSingleTreeBuilder` class is the single star-tree builder that uses off-heap memory. @@ -43,12 +44,22 @@ public class OffHeapSingleTreeBuilder extends BaseSingleTreeBuilder { private static final String STAR_TREE_RECORD_FILE_NAME = "star-tree.record"; // If the temporary buffer needed is larger than 500M, use MMAP, otherwise use DIRECT private static final long MMAP_SIZE_THRESHOLD = 500_000_000; + private static final byte[] EMPTY_BYTES = new byte[0]; private final File _segmentRecordFile; private final File _starTreeRecordFile; private final BufferedOutputStream _starTreeRecordOutputStream; private final RecordOffsets _starTreeRecordOffsets; + /// One bitmap per metric holding the doc ids whose group aggregated over no non-null input, or `null` when null + /// handling is disabled. A record is stored the way a regular column is: the metric keeps a placeholder in the + /// serialized record and its nullness lives beside it, so the record layout is unchanged and the arithmetic in + /// [FixedSizeRecordOffsets] still holds. + /// + /// Keying on the doc id is safe because records are never moved: sorting permutes an array of doc ids and compares + /// through it, leaving each record where it was written. + private final RoaringBitmap[] _metricNullBitmaps; + private PinotDataBuffer _starTreeRecordBuffer; private int _numReadableStarTreeRecords; @@ -64,6 +75,7 @@ public OffHeapSingleTreeBuilder(StarTreeV2BuilderConfig builderConfig, File outp _starTreeRecordFile); _starTreeRecordOutputStream = new BufferedOutputStream(new FileOutputStream(_starTreeRecordFile)); _starTreeRecordOffsets = createRecordOffsets(); + _metricNullBitmaps = _nullHandlingEnabled ? new RoaringBitmap[_numMetrics] : null; } /// Returns [FixedSizeRecordOffsets] when all metrics are serialized with a fixed size (see @@ -86,6 +98,11 @@ private RecordOffsets createRecordOffsets() { return new FixedSizeRecordOffsets(recordSize); } + /// Serializes a record into the temporary star-tree record store. + /// + /// A metric that aggregated over no non-null input has no value to write, so the aggregated type's zero is written + /// in its place and [#appendRecord] records the nullness in [#_metricNullBitmaps]. The placeholder is never read + /// back: [#deserializeStarTreeRecord] consults the bitmap first. @SuppressWarnings("unchecked") private byte[] serializeStarTreeRecord(Record starTreeRecord) { int numBytes = _numDimensions * Integer.BYTES; @@ -99,7 +116,8 @@ private byte[] serializeStarTreeRecord(Record starTreeRecord) { numBytes += Double.BYTES; break; case BYTES: - metricBytes[i] = _valueAggregators[i].serializeAggregatedValue(starTreeRecord._metrics[i]); + Object bytesValue = starTreeRecord._metrics[i]; + metricBytes[i] = bytesValue != null ? _valueAggregators[i].serializeAggregatedValue(bytesValue) : EMPTY_BYTES; numBytes += Integer.BYTES + metricBytes[i].length; break; default: @@ -114,10 +132,12 @@ private byte[] serializeStarTreeRecord(Record starTreeRecord) { for (int i = 0; i < _numMetrics; i++) { switch (_valueAggregators[i].getAggregatedValueType()) { case LONG: - byteBuffer.putLong((Long) starTreeRecord._metrics[i]); + Object longValue = starTreeRecord._metrics[i]; + byteBuffer.putLong(longValue != null ? (long) longValue : 0L); break; case DOUBLE: - byteBuffer.putDouble((Double) starTreeRecord._metrics[i]); + Object doubleValue = starTreeRecord._metrics[i]; + byteBuffer.putDouble(doubleValue != null ? (double) doubleValue : 0d); break; case BYTES: byteBuffer.putInt(metricBytes[i].length); @@ -130,7 +150,11 @@ private byte[] serializeStarTreeRecord(Record starTreeRecord) { return bytes; } - private Record deserializeStarTreeRecord(PinotDataBuffer buffer, long offset) { + /// Deserializes the record at `docId`, whose metrics start at `offset`. + /// + /// A metric marked null in [#_metricNullBitmaps] comes back as `null` without its placeholder being deserialized, + /// which matters for `BYTES`: the placeholder is an empty array that no aggregator can decode. + private Record deserializeStarTreeRecord(PinotDataBuffer buffer, long offset, int docId) { int[] dimensions = new int[_numDimensions]; for (int i = 0; i < _numDimensions; i++) { dimensions[i] = buffer.getInt(offset); @@ -140,20 +164,24 @@ private Record deserializeStarTreeRecord(PinotDataBuffer buffer, long offset) { for (int i = 0; i < _numMetrics; i++) { switch (_valueAggregators[i].getAggregatedValueType()) { case LONG: - metrics[i] = buffer.getLong(offset); + metrics[i] = hasMetricValue(docId, i) ? buffer.getLong(offset) : null; offset += Long.BYTES; break; case DOUBLE: - metrics[i] = buffer.getDouble(offset); + metrics[i] = hasMetricValue(docId, i) ? buffer.getDouble(offset) : null; offset += Double.BYTES; break; case BYTES: int numBytes = buffer.getInt(offset); offset += Integer.BYTES; - byte[] bytes = new byte[numBytes]; - buffer.copyTo(offset, bytes); + if (hasMetricValue(docId, i)) { + byte[] bytes = new byte[numBytes]; + buffer.copyTo(offset, bytes); + metrics[i] = _valueAggregators[i].deserializeAggregatedValue(bytes); + } else { + metrics[i] = null; + } offset += numBytes; - metrics[i] = _valueAggregators[i].deserializeAggregatedValue(bytes); break; default: throw new IllegalStateException(); @@ -168,13 +196,40 @@ void appendRecord(Record record) byte[] bytes = serializeStarTreeRecord(record); _starTreeRecordOutputStream.write(bytes); _starTreeRecordOffsets.addRecord(bytes.length); + if (_metricNullBitmaps != null) { + // The caller assigns this record _numDocs and increments it afterwards, so it is this record's doc id + markNullMetrics(record, _numDocs); + } + } + + /// Records which of the record's metrics aggregated over no non-null input. + private void markNullMetrics(Record record, int docId) { + for (int i = 0; i < _numMetrics; i++) { + if (record._metrics[i] == null) { + RoaringBitmap nullBitmap = _metricNullBitmaps[i]; + if (nullBitmap == null) { + nullBitmap = new RoaringBitmap(); + _metricNullBitmaps[i] = nullBitmap; + } + nullBitmap.add(docId); + } + } + } + + /// Returns whether the record holds a value for the metric, as opposed to having aggregated over no non-null input. + private boolean hasMetricValue(int docId, int metricId) { + if (_metricNullBitmaps == null) { + return true; + } + RoaringBitmap nullBitmap = _metricNullBitmaps[metricId]; + return nullBitmap == null || !nullBitmap.contains(docId); } @Override Record getStarTreeRecord(int docId) throws IOException { ensureBufferReadable(docId); - return deserializeStarTreeRecord(_starTreeRecordBuffer, _starTreeRecordOffsets.getStartOffset(docId)); + return deserializeStarTreeRecord(_starTreeRecordBuffer, _starTreeRecordOffsets.getStartOffset(docId), docId); } @Override diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/builder/StarTreeIndexCombiner.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/builder/StarTreeIndexCombiner.java index c3f4d685826f..64d0804ea5bb 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/builder/StarTreeIndexCombiner.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/builder/StarTreeIndexCombiner.java @@ -62,6 +62,7 @@ public List> combine(StarTreeV2BuilderConfig builderC File dimensionIndexFile = new File(starTreeIndexDir, dimension + V1Constants.Indexes.UNSORTED_SV_FORWARD_INDEX_FILE_EXTENSION); indexMap.add(Pair.of(new IndexKey(IndexType.FORWARD_INDEX, dimension), writeFile(dimensionIndexFile))); + writeNullValueVectorIfExists(starTreeIndexDir, dimension, indexMap); } // Write metric (function-column pair) indexes @@ -70,12 +71,27 @@ public List> combine(StarTreeV2BuilderConfig builderC File metricIndexFile = new File(starTreeIndexDir, metric + V1Constants.Indexes.RAW_SV_FORWARD_INDEX_FILE_EXTENSION); indexMap.add(Pair.of(new IndexKey(IndexType.FORWARD_INDEX, metric), writeFile(metricIndexFile))); + writeNullValueVectorIfExists(starTreeIndexDir, metric, indexMap); } FileUtils.cleanDirectory(starTreeIndexDir); return indexMap; } + /// Writes the null value vector for the given column if the builder created one. + /// + /// Only a null-aware star-tree creates them, and only for columns that actually contain null values, so a missing + /// file simply means there is nothing to record. + private void writeNullValueVectorIfExists(File starTreeIndexDir, String column, + List> indexMap) + throws IOException { + File nullValueVectorFile = + new File(starTreeIndexDir, column + V1Constants.Indexes.NULLVALUE_VECTOR_FILE_EXTENSION); + if (nullValueVectorFile.exists()) { + indexMap.add(Pair.of(new IndexKey(IndexType.NULL_VALUE_VECTOR, column), writeFile(nullValueVectorFile))); + } + } + private IndexValue writeFile(File srcFile) throws IOException { try (FileChannel src = new RandomAccessFile(srcFile, "r").getChannel()) { diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/builder/StarTreeIndexSeparator.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/builder/StarTreeIndexSeparator.java index 6e28093c9c96..f69505d11743 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/builder/StarTreeIndexSeparator.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/builder/StarTreeIndexSeparator.java @@ -112,6 +112,11 @@ private void separate(File starTreeOutputDir, int treeIndex) destIndexFile = new File(starTreeOutputDir, key._column + suffix); writeIndexToFile(destIndexFile, indexMap.get(key)); break; + case NULL_VALUE_VECTOR: + destIndexFile = + new File(starTreeOutputDir, key._column + V1Constants.Indexes.NULLVALUE_VECTOR_FILE_EXTENSION); + writeIndexToFile(destIndexFile, indexMap.get(key)); + break; default: } } diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/builder/StarTreeV2BuilderConfig.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/builder/StarTreeV2BuilderConfig.java index 765177c9acd4..d2455145e866 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/builder/StarTreeV2BuilderConfig.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/builder/StarTreeV2BuilderConfig.java @@ -55,8 +55,10 @@ public class StarTreeV2BuilderConfig { private final Set _skipStarNodeCreationForDimensions; private final TreeMap _aggregationSpecs; private final int _maxLeafRecords; + private final boolean _nullHandlingEnabled; public static StarTreeV2BuilderConfig fromIndexConfig(StarTreeIndexConfig indexConfig) { + boolean nullHandlingEnabled = indexConfig.isNullHandlingEnabled(); List dimensionsSplitOrder = indexConfig.getDimensionsSplitOrder(); Set skipStarNodeCreationForDimensions; @@ -72,7 +74,7 @@ public static StarTreeV2BuilderConfig fromIndexConfig(StarTreeIndexConfig indexC if (indexConfig.getFunctionColumnPairs() != null) { for (String functionColumnPair : indexConfig.getFunctionColumnPairs()) { AggregationFunctionColumnPair aggregationFunctionColumnPair = - AggregationFunctionColumnPair.fromColumnName(functionColumnPair); + AggregationFunctionColumnPair.fromColumnName(functionColumnPair, nullHandlingEnabled); AggregationFunctionColumnPair storedType = AggregationFunctionColumnPair.resolveToStoredType(aggregationFunctionColumnPair); // If there is already an equivalent functionColumnPair in the map, do not load another. @@ -83,7 +85,7 @@ public static StarTreeV2BuilderConfig fromIndexConfig(StarTreeIndexConfig indexC if (indexConfig.getAggregationConfigs() != null) { for (StarTreeAggregationConfig aggregationConfig : indexConfig.getAggregationConfigs()) { AggregationFunctionColumnPair aggregationFunctionColumnPair = - AggregationFunctionColumnPair.fromAggregationConfig(aggregationConfig); + AggregationFunctionColumnPair.fromAggregationConfig(aggregationConfig, nullHandlingEnabled); AggregationFunctionColumnPair storedType = AggregationFunctionColumnPair.resolveToStoredType(aggregationFunctionColumnPair); // If there is already an equivalent functionColumnPair in the map, do not load another. @@ -98,13 +100,13 @@ public static StarTreeV2BuilderConfig fromIndexConfig(StarTreeIndexConfig indexC } return new StarTreeV2BuilderConfig(dimensionsSplitOrder, skipStarNodeCreationForDimensions, aggregationSpecs, - maxLeafRecords); + maxLeafRecords, nullHandlingEnabled); } public static StarTreeV2BuilderConfig fromMetadata(StarTreeV2Metadata starTreeV2Metadata) { return new StarTreeV2BuilderConfig(starTreeV2Metadata.getDimensionsSplitOrder(), starTreeV2Metadata.getSkipStarNodeCreationForDimensions(), starTreeV2Metadata.getAggregationSpecs(), - starTreeV2Metadata.getMaxLeafRecords()); + starTreeV2Metadata.getMaxLeafRecords(), starTreeV2Metadata.isNullHandlingEnabled()); } /// Generates default config based on the segment metadata. @@ -173,8 +175,8 @@ public static StarTreeV2BuilderConfig generateDefaultConfig(SegmentMetadata segm AggregationSpec.DEFAULT); } - return new StarTreeV2BuilderConfig(dimensionsSplitOrder, Set.of(), aggregationSpecs, - DEFAULT_MAX_LEAF_RECORDS); + return new StarTreeV2BuilderConfig(dimensionsSplitOrder, Set.of(), aggregationSpecs, DEFAULT_MAX_LEAF_RECORDS, + false); } public static StarTreeV2BuilderConfig generateDefaultConfig(Schema schema, JsonNode columnsMetadata) { @@ -239,8 +241,8 @@ public static StarTreeV2BuilderConfig generateDefaultConfig(Schema schema, JsonN AggregationSpec.DEFAULT); } - return new StarTreeV2BuilderConfig(dimensionsSplitOrder, Set.of(), aggregationSpecs, - DEFAULT_MAX_LEAF_RECORDS); + return new StarTreeV2BuilderConfig(dimensionsSplitOrder, Set.of(), aggregationSpecs, DEFAULT_MAX_LEAF_RECORDS, + false); } public static Map convertJsonNodeToMap(JsonNode columnsMetadata) { @@ -253,11 +255,13 @@ public static Map convertJsonNodeToMap(JsonNode columnsMetadat } private StarTreeV2BuilderConfig(List dimensionsSplitOrder, Set skipStarNodeCreationForDimensions, - TreeMap aggregationSpecs, int maxLeafRecords) { + TreeMap aggregationSpecs, int maxLeafRecords, + boolean nullHandlingEnabled) { _dimensionsSplitOrder = dimensionsSplitOrder; _skipStarNodeCreationForDimensions = skipStarNodeCreationForDimensions; _aggregationSpecs = aggregationSpecs; _maxLeafRecords = maxLeafRecords; + _nullHandlingEnabled = nullHandlingEnabled; } public List getDimensionsSplitOrder() { @@ -280,10 +284,16 @@ public int getMaxLeafRecords() { return _maxLeafRecords; } + /// Returns whether the star-tree should be built with null-aware semantics. See + /// [StarTreeIndexConfig#isNullHandlingEnabled]. + public boolean isNullHandlingEnabled() { + return _nullHandlingEnabled; + } + /// Writes the metadata which is used to initialize the [StarTreeV2Metadata] when loading the segment. public void writeMetadata(Configuration metadataProperties, int totalDocs) { StarTreeV2Metadata.writeMetadata(metadataProperties, totalDocs, _dimensionsSplitOrder, _aggregationSpecs, - _maxLeafRecords, _skipStarNodeCreationForDimensions); + _maxLeafRecords, _skipStarNodeCreationForDimensions, _nullHandlingEnabled); } @Override @@ -295,20 +305,23 @@ public boolean equals(Object o) { return false; } StarTreeV2BuilderConfig that = (StarTreeV2BuilderConfig) o; - return _maxLeafRecords == that._maxLeafRecords && Objects.equals(_dimensionsSplitOrder, that._dimensionsSplitOrder) + return _maxLeafRecords == that._maxLeafRecords && _nullHandlingEnabled == that._nullHandlingEnabled + && Objects.equals(_dimensionsSplitOrder, that._dimensionsSplitOrder) && Objects.equals(_skipStarNodeCreationForDimensions, that._skipStarNodeCreationForDimensions) && Objects.equals(_aggregationSpecs, that._aggregationSpecs); } @Override public int hashCode() { - return Objects.hash(_dimensionsSplitOrder, _skipStarNodeCreationForDimensions, _aggregationSpecs, _maxLeafRecords); + return Objects.hash(_dimensionsSplitOrder, _skipStarNodeCreationForDimensions, _aggregationSpecs, _maxLeafRecords, + _nullHandlingEnabled); } @Override public String toString() { return new ToStringBuilder(this, ToStringStyle.SHORT_PREFIX_STYLE).append("splitOrder", _dimensionsSplitOrder) .append("skipStarNodeCreation", _skipStarNodeCreationForDimensions) - .append("aggregationSpecs", _aggregationSpecs).append("maxLeafRecords", _maxLeafRecords).toString(); + .append("aggregationSpecs", _aggregationSpecs).append("maxLeafRecords", _maxLeafRecords) + .append("nullHandlingEnabled", _nullHandlingEnabled).toString(); } } diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/store/StarTreeDataSource.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/store/StarTreeDataSource.java index 0b9df1286487..ac62636458c6 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/store/StarTreeDataSource.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/store/StarTreeDataSource.java @@ -26,19 +26,26 @@ import org.apache.pinot.segment.spi.index.column.ColumnIndexContainer; import org.apache.pinot.segment.spi.index.reader.Dictionary; import org.apache.pinot.segment.spi.index.reader.ForwardIndexReader; +import org.apache.pinot.segment.spi.index.reader.NullValueVectorReader; import org.apache.pinot.segment.spi.partition.PartitionFunction; import org.apache.pinot.spi.data.FieldSpec; public class StarTreeDataSource extends BaseDataSource { + /// Creates a data source over a star-tree column. + /// + /// `nullValueVector` is only non-null for a null-aware star-tree column that contains null values. Dimensions need + /// one just as metrics do: the reserved null dictionary id is out of range for the segment dictionary that the + /// star-tree shares, so the query side must recognize nulls from the vector rather than from the stored id. public StarTreeDataSource(FieldSpec fieldSpec, int numDocs, ForwardIndexReader forwardIndex, - @Nullable Dictionary dictionary) { + @Nullable Dictionary dictionary, @Nullable NullValueVectorReader nullValueVector) { super( new StarTreeDataSourceMetadata(fieldSpec, numDocs), new ColumnIndexContainer.FromMap.Builder() .with(StandardIndexes.forward(), forwardIndex) .with(StandardIndexes.dictionary(), dictionary) + .with(StandardIndexes.nullValueVector(), nullValueVector) .build()); } diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/store/StarTreeIndexMapUtils.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/store/StarTreeIndexMapUtils.java index 25217b881e1c..032793f58df8 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/store/StarTreeIndexMapUtils.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/store/StarTreeIndexMapUtils.java @@ -62,6 +62,9 @@ /// - 1.null.STAR_TREE.OFFSET = 5500 /// - 1.null.STAR_TREE.SIZE = 2500 /// - ... +/// +/// A null-aware star-tree additionally stores a NULL_VALUE_VECTOR entry per dimension and metric that has at least one +/// null value, e.g. `0.sum__metric.NULL_VALUE_VECTOR.OFFSET`. public class StarTreeIndexMapUtils { private StarTreeIndexMapUtils() { } @@ -75,7 +78,9 @@ private StarTreeIndexMapUtils() { /// Type of the index. public enum IndexType { - STAR_TREE, FORWARD_INDEX + // NOTE: NULL_VALUE_VECTOR is only written by a null-aware star-tree. A server running a version that predates it + // fails to parse the index map of such a segment, so it must not be enabled before the whole cluster is upgraded. + STAR_TREE, FORWARD_INDEX, NULL_VALUE_VECTOR } /// Key of the index map. @@ -189,10 +194,13 @@ public static List> loadFromInputStream(InputStream in } // Convert metric (function-column pair) to stored name for backward-compatibility if (!dimensionSet.contains(column)) { - AggregationFunctionColumnPair functionColumnPair = AggregationFunctionColumnPair.fromColumnName(column); + // A count__column pair is only ever stored by a null-aware star-tree, where it holds the count of the + // non-null values of that column rather than the row count + AggregationFunctionColumnPair functionColumnPair = AggregationFunctionColumnPair.fromColumnName(column, + starTreeMetadataList.get(starTreeId).isNullHandlingEnabled()); column = AggregationFunctionColumnPair.resolveToStoredType(functionColumnPair).toColumnName(); } - indexKey = new IndexKey(IndexType.FORWARD_INDEX, column); + indexKey = new IndexKey(indexType, column); } long value = configuration.getLong(key); diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/store/StarTreeLoaderUtils.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/store/StarTreeLoaderUtils.java index 4e6ab15b49bb..d2ebd01e763a 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/store/StarTreeLoaderUtils.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/store/StarTreeLoaderUtils.java @@ -23,8 +23,11 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import javax.annotation.Nullable; import org.apache.pinot.segment.local.aggregator.ValueAggregatorFactory; +import org.apache.pinot.segment.local.io.util.PinotDataBitSet; import org.apache.pinot.segment.local.segment.index.forward.ForwardIndexReaderFactory; +import org.apache.pinot.segment.local.segment.index.readers.NullValueVectorReaderImpl; import org.apache.pinot.segment.local.segment.index.readers.forward.FixedBitSVForwardIndexReaderV2; import org.apache.pinot.segment.local.startree.OffHeapStarTree; import org.apache.pinot.segment.spi.ColumnMetadata; @@ -33,6 +36,7 @@ import org.apache.pinot.segment.spi.index.column.ColumnIndexContainer; import org.apache.pinot.segment.spi.index.metadata.SegmentMetadataImpl; import org.apache.pinot.segment.spi.index.reader.ForwardIndexReader; +import org.apache.pinot.segment.spi.index.reader.NullValueVectorReader; import org.apache.pinot.segment.spi.index.startree.AggregationFunctionColumnPair; import org.apache.pinot.segment.spi.index.startree.StarTree; import org.apache.pinot.segment.spi.index.startree.StarTreeV2; @@ -63,16 +67,23 @@ public static List loadStarTreeV2(SegmentDirectory.Reader segmentRea StarTreeV2Metadata starTreeMetadata = starTreeMetadataList.get(i); int numDocs = starTreeMetadata.getNumDocs(); + boolean nullHandlingEnabled = starTreeMetadata.isNullHandlingEnabled(); Map dataSourceMap = new HashMap<>(); // Load dimension forward indexes for (String dimension : starTreeMetadata.getDimensionsSplitOrder()) { PinotDataBuffer forwardIndexDataBuffer = indexReader.getIndexFor(dimension, StandardIndexes.forward()); ColumnMetadata columnMetadata = segmentMetadata.getColumnMetadataFor(dimension); + // A null-aware star-tree reserves one dictionary id per dimension for null values, so its forward index is + // created with cardinality + 1 values and may be one bit wider than the segment column's own forward index + int numBitsPerElement = nullHandlingEnabled + ? PinotDataBitSet.getNumBitsPerValue(columnMetadata.getCardinality()) + : columnMetadata.getBitsPerElement(); FixedBitSVForwardIndexReaderV2 forwardIndex = - new FixedBitSVForwardIndexReaderV2(forwardIndexDataBuffer, numDocs, columnMetadata.getBitsPerElement()); + new FixedBitSVForwardIndexReaderV2(forwardIndexDataBuffer, numDocs, numBitsPerElement); dataSourceMap.put(dimension, new StarTreeDataSource(columnMetadata.getFieldSpec(), numDocs, forwardIndex, - indexContainerMap.get(dimension).getIndex(StandardIndexes.dictionary()))); + indexContainerMap.get(dimension).getIndex(StandardIndexes.dictionary()), + loadNullValueVector(indexReader, dimension, nullHandlingEnabled))); } // Load metric (function-column pair) forward indexes @@ -83,7 +94,8 @@ public static List loadStarTreeV2(SegmentDirectory.Reader segmentRea FieldSpec fieldSpec = new MetricFieldSpec(metric, dataType); ForwardIndexReader forwardIndex = ForwardIndexReaderFactory.getInstance() .createRawIndexReader(forwardIndexDataBuffer, dataType.getStoredType(), true); - dataSourceMap.put(metric, new StarTreeDataSource(fieldSpec, numDocs, forwardIndex, null)); + dataSourceMap.put(metric, new StarTreeDataSource(fieldSpec, numDocs, forwardIndex, null, + loadNullValueVector(indexReader, metric, nullHandlingEnabled))); } starTrees.add(new StarTreeV2() { @@ -115,4 +127,18 @@ public void close() } return starTrees; } + + /// Returns the null value vector of a null-aware star-tree column, or `null` when the column has none. + /// + /// A regular star-tree never stores one, and a null-aware star-tree only stores one for columns that actually + /// contain null values. + @Nullable + private static NullValueVectorReader loadNullValueVector(SegmentDirectory.Reader indexReader, String column, + boolean nullHandlingEnabled) + throws IOException { + if (!nullHandlingEnabled || !indexReader.hasIndexFor(column, StandardIndexes.nullValueVector())) { + return null; + } + return new NullValueVectorReaderImpl(indexReader.getIndexFor(column, StandardIndexes.nullValueVector())); + } } diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/TableConfigUtils.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/TableConfigUtils.java index 8ef3c48727e0..2ece7905724f 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/TableConfigUtils.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/TableConfigUtils.java @@ -2000,11 +2000,14 @@ private static void validateStarTreeIndexConfigs(List starT Set functionColumnPairsSet = new HashSet<>(); Set storedTypes = new HashSet<>(); Set aggregatedColumns = new HashSet<>(); + // A null-aware star-tree stores COUNT per column rather than collapsing every COUNT to COUNT(*), so the column + // of a COUNT pair has to be kept and validated against the schema + boolean nullHandlingEnabled = starTreeIndexConfig.isNullHandlingEnabled(); if (functionColumnPairs != null) { for (String functionColumnPair : functionColumnPairs) { AggregationFunctionColumnPair columnPair; try { - columnPair = AggregationFunctionColumnPair.fromColumnName(functionColumnPair); + columnPair = AggregationFunctionColumnPair.fromColumnName(functionColumnPair, nullHandlingEnabled); } catch (Exception e) { throw new IllegalStateException("Invalid StarTreeIndex config: " + functionColumnPair + ". Must be" + "in the form __"); @@ -2032,7 +2035,7 @@ private static void validateStarTreeIndexConfigs(List starT for (StarTreeAggregationConfig aggregationConfig : aggregationConfigs) { AggregationFunctionColumnPair columnPair; try { - columnPair = AggregationFunctionColumnPair.fromAggregationConfig(aggregationConfig); + columnPair = AggregationFunctionColumnPair.fromAggregationConfig(aggregationConfig, nullHandlingEnabled); } catch (Exception e) { throw new IllegalStateException("Invalid StarTreeIndex config: " + aggregationConfig); } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/startree/v2/builder/NullAwareStarTreeBuilderTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/startree/v2/builder/NullAwareStarTreeBuilderTest.java new file mode 100644 index 000000000000..11dd454ad0fb --- /dev/null +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/startree/v2/builder/NullAwareStarTreeBuilderTest.java @@ -0,0 +1,230 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.segment.local.startree.v2.builder; + +import java.io.File; +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; +import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; +import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader; +import org.apache.pinot.segment.local.startree.v2.builder.MultipleTreesBuilder.BuildMode; +import org.apache.pinot.segment.spi.AggregationFunctionType; +import org.apache.pinot.segment.spi.ImmutableSegment; +import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; +import org.apache.pinot.segment.spi.index.reader.ForwardIndexReader; +import org.apache.pinot.segment.spi.index.reader.ForwardIndexReaderContext; +import org.apache.pinot.segment.spi.index.reader.NullValueVectorReader; +import org.apache.pinot.segment.spi.index.startree.AggregationFunctionColumnPair; +import org.apache.pinot.segment.spi.index.startree.StarTreeV2; +import org.apache.pinot.spi.config.table.StarTreeIndexConfig; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.data.readers.GenericRow; +import org.apache.pinot.spi.utils.ReadMode; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.roaringbitmap.buffer.ImmutableRoaringBitmap; +import org.testng.annotations.AfterMethod; +import org.testng.annotations.DataProvider; +import org.testng.annotations.Test; + +import static org.testng.Assert.*; + + +/// Builds star-trees over a segment whose metric is null for a whole group. +/// +/// The group with no non-null input is the case a null-aware star-tree exists to represent, and it is the case the +/// regular star-tree cannot distinguish from one that genuinely aggregated the column's default null value. +/// +/// The build mode is a parameter rather than a random choice, because the two builders represent an all-null group +/// differently: the on-heap one keeps the aggregated value in memory where `null` needs no encoding, while the +/// off-heap one serializes every record to a temporary store and has to carry the nullness beside it. +public class NullAwareStarTreeBuilderTest { + private static final File TEMP_DIR = new File(FileUtils.getTempDirectory(), "NullAwareStarTreeBuilderTest"); + private static final String TABLE_NAME = "testTable"; + private static final String SEGMENT_NAME = "testSegment"; + private static final String DIMENSION = "d"; + private static final String METRIC = "m"; + + /// One record per leaf, so every parent node has to aggregate rather than pass a single record through. + private static final int MAX_LEAF_RECORDS = 1; + + /// `d = 0` holds only null metrics, `d = 1` holds only non-null ones. The default null value of an `INT` metric is + /// `0`, which is below both real values, so a regular star-tree reports `0` as the minimum where a null-aware one + /// reports `10`. + private static final int[] DIMENSION_VALUES = {0, 0, 1, 1}; + private static final Integer[] METRIC_VALUES = {null, null, 10, 20}; + + private static final String MIN_COLUMN = + new AggregationFunctionColumnPair(AggregationFunctionType.MIN, METRIC).toColumnName(); + + @AfterMethod + public void cleanUp() + throws IOException { + FileUtils.deleteDirectory(TEMP_DIR); + } + + @DataProvider(name = "buildModes") + public Object[][] buildModes() { + return new Object[][]{{BuildMode.ON_HEAP}, {BuildMode.OFF_HEAP}}; + } + + /// A group with no non-null input has no aggregated value to store, and the off-heap builder used to have nowhere + /// to put that: it serialized every metric unconditionally and dereferenced the missing value. + @Test(dataProvider = "buildModes") + public void anAllNullGroupSurvivesTheBuild(BuildMode buildMode) + throws Exception { + File indexDir = createSegment(); + buildStarTrees(indexDir, buildMode, starTreeConfig(true)); + + ImmutableSegment segment = ImmutableSegmentLoader.load(indexDir, ReadMode.mmap); + try { + StarTreeV2 starTree = segment.getStarTrees().get(0); + assertTrue(starTree.getMetadata().isNullHandlingEnabled()); + + // The all-null group is recorded in the metric's null vector rather than as an aggregated value + ImmutableRoaringBitmap nullBitmap = nullBitmap(starTree, MIN_COLUMN); + assertNotNull(nullBitmap, "A null-aware star-tree must write a null vector for a metric with an all-null group"); + assertFalse(nullBitmap.isEmpty()); + + // Every group that did aggregate something reports a real minimum, never the column's default null value + List minimums = nonNullValues(starTree, MIN_COLUMN, nullBitmap); + assertFalse(minimums.isEmpty(), "Some group must have aggregated a value, or the check below proves nothing"); + for (double minimum : minimums) { + assertEquals(minimum, 10.0, "A null row must not be aggregated as the column default"); + } + } finally { + segment.destroy(); + } + } + + /// The two variants answer differently, so a segment has to be able to hold both: the builder config's identity + /// includes the flag, which is what stops one being reused for the other. + @Test(dataProvider = "buildModes") + public void bothVariantsCoexistAndDisagreeOnNulls(BuildMode buildMode) + throws Exception { + File indexDir = createSegment(); + buildStarTrees(indexDir, buildMode, starTreeConfig(false), starTreeConfig(true)); + + ImmutableSegment segment = ImmutableSegmentLoader.load(indexDir, ReadMode.mmap); + try { + List starTrees = segment.getStarTrees(); + assertEquals(starTrees.size(), 2, "The two variants differ only by the flag, and must not be deduplicated"); + + StarTreeV2 regular = variant(starTrees, false); + StarTreeV2 nullAware = variant(starTrees, true); + + // Same shape: the flag is the only difference between the two configs + assertEquals(regular.getMetadata().getDimensionsSplitOrder(), nullAware.getMetadata().getDimensionsSplitOrder()); + assertEquals(regular.getMetadata().getFunctionColumnPairs(), nullAware.getMetadata().getFunctionColumnPairs()); + + // The regular tree folds nulls into the column default, so it has no null vector and reports that default + assertNull(nullBitmap(regular, MIN_COLUMN), "A regular star-tree does not write a null vector"); + assertTrue(nonNullValues(regular, MIN_COLUMN, null).contains(0.0), + "A regular star-tree aggregates a null row as the column's default null value"); + + // The null-aware tree excludes them, so the default never appears as a minimum + ImmutableRoaringBitmap nullBitmap = nullBitmap(nullAware, MIN_COLUMN); + assertNotNull(nullBitmap); + assertFalse(nonNullValues(nullAware, MIN_COLUMN, nullBitmap).contains(0.0), + "A null-aware star-tree excludes null rows from the pre-aggregation"); + } finally { + segment.destroy(); + } + } + + private static StarTreeV2 variant(List starTrees, boolean nullHandlingEnabled) { + for (StarTreeV2 starTree : starTrees) { + if (starTree.getMetadata().isNullHandlingEnabled() == nullHandlingEnabled) { + return starTree; + } + } + throw new AssertionError("No star-tree with nullHandlingEnabled=" + nullHandlingEnabled); + } + + private static ImmutableRoaringBitmap nullBitmap(StarTreeV2 starTree, String column) { + NullValueVectorReader nullValueVector = starTree.getDataSource(column).getNullValueVector(); + return nullValueVector != null ? nullValueVector.getNullBitmap() : null; + } + + /// Returns the pre-aggregated values of every doc the null vector does not mark, so that a placeholder left behind + /// for a null group is never read as data. + private static List nonNullValues(StarTreeV2 starTree, String column, ImmutableRoaringBitmap nullBitmap) + throws IOException { + List values = new ArrayList<>(); + ForwardIndexReader reader = starTree.getDataSource(column).getForwardIndex(); + assertNotNull(reader); + try (ForwardIndexReaderContext context = reader.createContext()) { + for (int docId = 0; docId < starTree.getMetadata().getNumDocs(); docId++) { + if (nullBitmap == null || !nullBitmap.contains(docId)) { + values.add(reader.getDouble(docId, context)); + } + } + } + return values; + } + + private static StarTreeIndexConfig starTreeConfig(boolean nullHandlingEnabled) { + return new StarTreeIndexConfig(List.of(DIMENSION), null, + List.of(new AggregationFunctionColumnPair(AggregationFunctionType.MIN, METRIC).toColumnName()), null, + MAX_LEAF_RECORDS, nullHandlingEnabled); + } + + private static void buildStarTrees(File indexDir, BuildMode buildMode, StarTreeIndexConfig... configs) + throws Exception { + try (MultipleTreesBuilder builder = new MultipleTreesBuilder(List.of(configs), false, indexDir, buildMode)) { + builder.build(); + } + } + + private static File createSegment() + throws Exception { + Schema schema = new Schema.SchemaBuilder().setSchemaName(TABLE_NAME) + .addSingleValueDimension(DIMENSION, DataType.INT) + .addMetric(METRIC, DataType.INT) + .build(); + TableConfig tableConfig = + new TableConfigBuilder(TableType.OFFLINE).setTableName(TABLE_NAME).setNullHandlingEnabled(true).build(); + + List rows = new ArrayList<>(DIMENSION_VALUES.length); + for (int i = 0; i < DIMENSION_VALUES.length; i++) { + GenericRow row = new GenericRow(); + row.putValue(DIMENSION, DIMENSION_VALUES[i]); + if (METRIC_VALUES[i] != null) { + row.putValue(METRIC, METRIC_VALUES[i]); + } else { + // Ingestion stores the column default and records the row in the null vector + row.putDefaultNullValue(METRIC, schema.getFieldSpecFor(METRIC).getDefaultNullValue()); + } + rows.add(row); + } + + SegmentGeneratorConfig segmentGeneratorConfig = new SegmentGeneratorConfig(tableConfig, schema); + segmentGeneratorConfig.setOutDir(TEMP_DIR.getPath()); + segmentGeneratorConfig.setSegmentName(SEGMENT_NAME); + SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl(); + driver.init(segmentGeneratorConfig, new GenericRowRecordReader(rows)); + driver.build(); + return new File(TEMP_DIR, SEGMENT_NAME); + } +} diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/startree/v2/builder/StarTreeBuilderUtilsTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/startree/v2/builder/StarTreeBuilderUtilsTest.java index c9480238963f..8a819f2e2cbd 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/startree/v2/builder/StarTreeBuilderUtilsTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/startree/v2/builder/StarTreeBuilderUtilsTest.java @@ -214,7 +214,7 @@ public void testShouldModifyExistingStarTrees() { TreeMap aggregationSpecs = new TreeMap<>(); aggregationSpecs.put(new AggregationFunctionColumnPair(AggregationFunctionType.DISTINCTCOUNTHLL, "col2"), new AggregationSpec(null, null, null, null, null, Map.of(Constants.HLL_LOG2M_KEY, 16))); - StarTreeV2Metadata.writeMetadata(metadataProperties, 1, List.of("col1"), aggregationSpecs, 100, Set.of()); + StarTreeV2Metadata.writeMetadata(metadataProperties, 1, List.of("col1"), aggregationSpecs, 100, Set.of(), false); StarTreeV2Metadata existingStarTreeMetadata = new StarTreeV2Metadata(metadataProperties); StarTreeIndexConfig starTreeIndexConfig = new StarTreeIndexConfig(List.of("col1"), null, null, List.of( diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/startree/AggregationFunctionColumnPair.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/startree/AggregationFunctionColumnPair.java index 8c5a3347f48f..60801ab5bab6 100644 --- a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/startree/AggregationFunctionColumnPair.java +++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/startree/AggregationFunctionColumnPair.java @@ -34,14 +34,29 @@ public class AggregationFunctionColumnPair implements Comparable _skipStarNodeCreationForDimensions; + private final boolean _nullHandlingEnabled; public StarTreeV2Metadata(Configuration metadataProperties) { _numDocs = metadataProperties.getInt(MetadataKey.TOTAL_DOCS); _dimensionsSplitOrder = Arrays.asList(metadataProperties.getStringArray(MetadataKey.DIMENSIONS_SPLIT_ORDER)); + // NOTE: Must be read before the aggregation specs, which are parsed differently for a null-aware star-tree + _nullHandlingEnabled = metadataProperties.getBoolean(MetadataKey.NULL_HANDLING_ENABLED, false); _aggregationSpecs = new TreeMap<>(); int numAggregations = metadataProperties.getInt(MetadataKey.AGGREGATION_COUNT, 0); if (numAggregations > 0) { @@ -54,7 +57,10 @@ public StarTreeV2Metadata(Configuration metadataProperties) { AggregationFunctionType functionType = AggregationFunctionType.getAggregationFunctionType(aggregationConfig.getString(MetadataKey.FUNCTION_TYPE)); String columnName = aggregationConfig.getString(MetadataKey.COLUMN_NAME); - AggregationFunctionColumnPair functionColumnPair = new AggregationFunctionColumnPair(functionType, columnName); + AggregationFunctionColumnPair functionColumnPair = + functionType == AggregationFunctionType.COUNT && _nullHandlingEnabled + ? AggregationFunctionColumnPair.countColumn(columnName) + : new AggregationFunctionColumnPair(functionType, columnName); // Lookup the stored aggregation type AggregationFunctionColumnPair storedType = AggregationFunctionColumnPair.resolveToStoredType(functionColumnPair); @@ -79,7 +85,7 @@ public StarTreeV2Metadata(Configuration metadataProperties) { // Backward compatibility with columnName format for (String functionColumnPairName : metadataProperties.getStringArray(MetadataKey.FUNCTION_COLUMN_PAIRS)) { AggregationFunctionColumnPair functionColumnPair = - AggregationFunctionColumnPair.fromColumnName(functionColumnPairName); + AggregationFunctionColumnPair.fromColumnName(functionColumnPairName, _nullHandlingEnabled); // Lookup the stored aggregation type AggregationFunctionColumnPair storedType = AggregationFunctionColumnPair.resolveToStoredType(functionColumnPair); @@ -120,9 +126,17 @@ public Set getSkipStarNodeCreationForDimensions() { return _skipStarNodeCreationForDimensions; } + /// Returns whether this star-tree was pre-aggregated with null-aware semantics. + /// + /// A null-aware star-tree stores null dimension values under a dedicated dictionary id and excludes null metric + /// values from the pre-aggregation, so it is only consistent with queries that have null handling enabled. + public boolean isNullHandlingEnabled() { + return _nullHandlingEnabled; + } + public static void writeMetadata(Configuration metadataProperties, int totalDocs, List dimensionsSplitOrder, TreeMap aggregationSpecs, int maxLeafRecords, - Set skipStarNodeCreationForDimensions) { + Set skipStarNodeCreationForDimensions, boolean nullHandlingEnabled) { metadataProperties.setProperty(MetadataKey.TOTAL_DOCS, totalDocs); metadataProperties.setProperty(MetadataKey.DIMENSIONS_SPLIT_ORDER, dimensionsSplitOrder); metadataProperties.setProperty(MetadataKey.FUNCTION_COLUMN_PAIRS, aggregationSpecs.keySet()); @@ -154,5 +168,6 @@ public static void writeMetadata(Configuration metadataProperties, int totalDocs metadataProperties.setProperty(MetadataKey.MAX_LEAF_RECORDS, maxLeafRecords); metadataProperties.setProperty(MetadataKey.SKIP_STAR_NODE_CREATION_FOR_DIMENSIONS, skipStarNodeCreationForDimensions); + metadataProperties.setProperty(MetadataKey.NULL_HANDLING_ENABLED, nullHandlingEnabled); } } diff --git a/pinot-segment-spi/src/test/java/org/apache/pinot/segment/spi/index/startree/StarTreeV2MetadataTest.java b/pinot-segment-spi/src/test/java/org/apache/pinot/segment/spi/index/startree/StarTreeV2MetadataTest.java index e29d73abc2f2..1a68afaee628 100644 --- a/pinot-segment-spi/src/test/java/org/apache/pinot/segment/spi/index/startree/StarTreeV2MetadataTest.java +++ b/pinot-segment-spi/src/test/java/org/apache/pinot/segment/spi/index/startree/StarTreeV2MetadataTest.java @@ -98,7 +98,8 @@ public void testDuplicateFunctionColumnPairs() { private static Configuration createMetadata(List dimensionsSplitOrder, TreeMap aggregationSpecs) { Configuration metadataProperties = new PropertiesConfiguration(); - StarTreeV2Metadata.writeMetadata(metadataProperties, 1, dimensionsSplitOrder, aggregationSpecs, 10000, Set.of()); + StarTreeV2Metadata.writeMetadata(metadataProperties, 1, dimensionsSplitOrder, aggregationSpecs, 10000, Set.of(), + false); return metadataProperties; } diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/config/table/StarTreeIndexConfig.java b/pinot-spi/src/main/java/org/apache/pinot/spi/config/table/StarTreeIndexConfig.java index 542ee055c6e2..1e11483f5cdb 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/config/table/StarTreeIndexConfig.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/config/table/StarTreeIndexConfig.java @@ -38,6 +38,8 @@ public class StarTreeIndexConfig extends BaseJsonConfig { private final List _aggregationConfigs; // The upper bound of records to be scanned at the leaf node private final int _maxLeafRecords; + // Whether the star-tree pre-aggregates with null-aware semantics + private final boolean _nullHandlingEnabled; @JsonCreator public StarTreeIndexConfig( @@ -46,7 +48,8 @@ public StarTreeIndexConfig( List skipStarNodeCreationForDimensions, @JsonProperty(value = "functionColumnPairs") @Nullable List functionColumnPairs, @JsonProperty(value = "aggregationConfigs") @Nullable List aggregationConfigs, - @JsonProperty(value = "maxLeafRecords") int maxLeafRecords) { + @JsonProperty(value = "maxLeafRecords") int maxLeafRecords, + @JsonProperty(value = "nullHandlingEnabled") boolean nullHandlingEnabled) { Preconditions.checkArgument(CollectionUtils.isNotEmpty(dimensionsSplitOrder), "'dimensionsSplitOrder' must be configured"); _dimensionsSplitOrder = dimensionsSplitOrder; @@ -55,10 +58,20 @@ public StarTreeIndexConfig( _functionColumnPairs = CollectionUtils.isNotEmpty(functionColumnPairs) ? functionColumnPairs : null; _aggregationConfigs = CollectionUtils.isNotEmpty(aggregationConfigs) ? aggregationConfigs : null; _maxLeafRecords = maxLeafRecords; + _nullHandlingEnabled = nullHandlingEnabled; Preconditions.checkArgument(_functionColumnPairs != null || _aggregationConfigs != null, "Either 'functionColumnPairs' or 'aggregationConfigs' must be configured"); } + /// Convenience constructor for a star-tree that is not null-aware, matching the behavior before + /// `nullHandlingEnabled` was introduced. + public StarTreeIndexConfig(List dimensionsSplitOrder, + @Nullable List skipStarNodeCreationForDimensions, @Nullable List functionColumnPairs, + @Nullable List aggregationConfigs, int maxLeafRecords) { + this(dimensionsSplitOrder, skipStarNodeCreationForDimensions, functionColumnPairs, aggregationConfigs, + maxLeafRecords, false); + } + public List getDimensionsSplitOrder() { return _dimensionsSplitOrder; } @@ -81,4 +94,14 @@ public List getAggregationConfigs() { public int getMaxLeafRecords() { return _maxLeafRecords; } + + /// Returns whether this star-tree pre-aggregates with null-aware semantics. + /// + /// A null-aware star-tree keeps null dimension values in their own group instead of folding them into the column's + /// default null value, and excludes null metric values from the pre-aggregation. It can therefore only serve queries + /// with null handling enabled, while a regular star-tree can only serve queries with null handling disabled (or + /// enabled queries over columns that happen to contain no nulls). + public boolean isNullHandlingEnabled() { + return _nullHandlingEnabled; + } }