yashmayya commented on code in PR #18334:
URL: https://github.com/apache/pinot/pull/18334#discussion_r3632799718
##########
pinot-core/src/test/java/org/apache/pinot/core/plan/maker/MetadataAndDictionaryAggregationPlanMakerTest.java:
##########
@@ -170,6 +178,42 @@ public void testPlanMaker(String query, Class<? extends
Operator<?>> operatorCla
assertTrue(upsertOperatorClass.isInstance(upsertOperator));
}
+ /**
+ * Verifies the partial metadata-based aggregation path. When a query mixes
a metadata-eligible function (MAX) with a
+ * non-eligible one (SUM), the plan uses an {@link AggregationOperator} that
pre-aggregates the eligible function from
+ * metadata while scanning the rest. Before this feature, a non-scan based
operator was used only when <em>all</em>
+ * functions were metadata eligible; a mixed query would have scanned every
function.
+ * <p>
+ * To prove the eligible function is actually served from metadata (and not
scanned), the column dictionary is
+ * overridden to report a bogus max value that does not exist in the data.
The query result equals the bogus value
+ * only if the metadata path is taken; a scan would return the true max.
+ */
+ @Test
+ public void testPartialMetadataBasedAggregationServesEligibleFromMetadata() {
+ int bogusMax = 999_999_999;
+ QueryContext queryContext =
+ QueryContextConverterUtils.getQueryContext("select
max(daysSinceEpoch), sum(column1) from testTable");
+
+ // Override only daysSinceEpoch's dictionary max value; everything else
delegates to the real segment.
+ DataSource realDataSource = _indexSegment.getDataSource("daysSinceEpoch",
queryContext.getSchema());
+ Dictionary dictionaryWithBogusMax = mock(Dictionary.class,
delegatesTo(realDataSource.getDictionary()));
+ doReturn(bogusMax).when(dictionaryWithBogusMax).getMaxVal();
+ DataSource dataSourceWithBogusMax = mock(DataSource.class,
delegatesTo(realDataSource));
+
doReturn(dictionaryWithBogusMax).when(dataSourceWithBogusMax).getDictionary();
+ IndexSegment segment = mock(IndexSegment.class,
delegatesTo(_indexSegment));
+
doReturn(dataSourceWithBogusMax).when(segment).getDataSource(eq("daysSinceEpoch"),
any());
+
+ Operator<?> operator = PLAN_MAKER.makeSegmentPlanNode(new
SegmentContext(segment), queryContext).run();
+ // A mixed query must use the (partial) AggregationOperator, not the fully
non-scan based operator.
+ assertTrue(operator instanceof AggregationOperator);
+
+ AggregationResultsBlock resultsBlock = (AggregationResultsBlock)
operator.nextBlock();
+ List<Object> results = resultsBlock.getResults();
+ assertNotNull(results);
+ // MAX is served from the (overridden) dictionary metadata, so the bogus
value proves the metadata path was used.
+ assertEquals(((Number) results.get(0)).doubleValue(), (double) bogusMax);
Review Comment:
Nice trick with the bogus max to prove the metadata path was taken. A few
additions would make this airtight:
- also assert `sum(column1)` (`results.get(1)`) equals the real scanned sum,
so we know the scanned side still runs correctly alongside the resolved one.
- add regression cases for the two things this refactor specifically fixes:
a single non-resolvable agg with no filter (`select sum(column1)`) should not
throw and should fall to the scan path, and an agg over an expression arg
(`max(add(column1, column2))`) should be treated as non-resolvable.
- the numeric-type guard (MIN/MAX on a non-numeric column falling back to
scan) is the correctness fix in this PR but isn't covered anywhere - worth a
case here.
##########
pinot-core/src/main/java/org/apache/pinot/core/plan/AggregationPlanNode.java:
##########
@@ -115,17 +122,34 @@ public Operator<AggregationResultsBlock>
buildNonFilteredAggOperator() {
boolean hasNullValues = _queryContext.isNullHandlingEnabled() &&
hasNullValues(aggregationFunctions);
if (!hasNullValues) {
- // Priority 2: Check if non-scan based aggregation is feasible
- if (filterOperator.isResultMatchingAll() && isFitForNonScanBasedPlan()) {
+ // when the filter matches all documents, resolve as many functions as
possible from the column
+ // dictionary/metadata without scanning the segment. Eligibility is
evaluated once per function here
+ // and reused for both the fully non-scan path (all functions
resolvable) and
+ // the partial path (some functions resolvable).
+ if (filterOperator.isResultMatchingAll()) {
+ boolean[] metadataResolvable = new
boolean[aggregationFunctions.length];
DataSource[] dataSources = new DataSource[aggregationFunctions.length];
+ int numResolved = 0;
for (int i = 0; i < aggregationFunctions.length; i++) {
- List<?> inputExpressions =
aggregationFunctions[i].getInputExpressions();
- if (!inputExpressions.isEmpty()) {
- String column = ((ExpressionContext)
inputExpressions.get(0)).getIdentifier();
- dataSources[i] = _indexSegment.getDataSource(column,
_queryContext.getSchema());
+ DataSource dataSource =
getDataSourceForAggregationFunction(aggregationFunctions[i]);
+ if (isFitForNonScanBasedPlan(aggregationFunctions[i], dataSource)) {
+ metadataResolvable[i] = true;
+ dataSources[i] = dataSource;
+ numResolved++;
}
}
- return new NonScanBasedAggregationOperator(_queryContext, dataSources,
numTotalDocs);
+
+ if (numResolved == aggregationFunctions.length) {
+ // Priority 2: all functions can be resolved from
dictionary/metadata -> fully non-scan based execution
+ return new NonScanBasedAggregationOperator(_queryContext,
dataSources, numTotalDocs);
+ }
+ if (numResolved > 0) {
+ // some functions can be resolved from dictionary/metadata; the rest
fall back to scan-based
+ // execution in the AggregationOperator.
+ aggregationInfo =
AggregationFunctionUtils.buildAggregationInfoWithoutStarTree(_segmentContext,
_queryContext,
Review Comment:
Minor/optional: the partial path still projects every function's columns via
`buildAggregationInfoWithoutStarTree`, so a column used only by a
metadata-resolved function is still read during the scan - we save the per-row
`aggregate()` but not the column I/O. Usually the per-row compute is the bulk
(especially DISTINCT*), so this is fine, but if you wanted the full win you'd
need to exclude resolved-only columns from the projected set. OK to leave as a
follow-up.
##########
pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtilsTest.java:
##########
@@ -0,0 +1,86 @@
+/**
+ * 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.core.query.aggregation.function;
+
+import org.apache.pinot.segment.spi.AggregationFunctionType;
+import org.apache.pinot.segment.spi.datasource.DataSource;
+import org.apache.pinot.segment.spi.index.reader.Dictionary;
+import org.testng.annotations.Test;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertThrows;
+
+
+/**
+ * Unit test for {@link AggregationFunctionUtils#getAggregationResult}, the
metadata/dictionary based aggregation
+ * result resolver used by the non-scan based and partial metadata based
aggregation paths.
+ */
+@SuppressWarnings("rawtypes")
+public class AggregationFunctionUtilsTest {
+
+ private static AggregationFunction mockFunction(AggregationFunctionType
type) {
+ AggregationFunction aggregationFunction = mock(AggregationFunction.class);
+ when(aggregationFunction.getType()).thenReturn(type);
+ return aggregationFunction;
+ }
+
+ @Test
+ public void testCountResolvedFromNumTotalDocs() {
+ AggregationFunction countFunction =
mockFunction(AggregationFunctionType.COUNT);
+ // COUNT is resolved directly from numTotalDocs and must not touch the
(possibly null) data source.
+ Object result =
AggregationFunctionUtils.getAggregationResult(countFunction, null, 42, "TEST");
+ assertEquals(result, 42L);
+ }
+
+ @Test
+ public void testMinAndMaxResolvedFromDictionary() {
+ Dictionary dictionary = mock(Dictionary.class);
+ when(dictionary.getMinVal()).thenReturn(5);
+ when(dictionary.getMaxVal()).thenReturn(10);
+ DataSource dataSource = mock(DataSource.class);
+ when(dataSource.getDictionary()).thenReturn(dictionary);
+
+ Object minResult =
AggregationFunctionUtils.getAggregationResult(mockFunction(AggregationFunctionType.MIN),
+ dataSource, 100, "TEST");
+ assertEquals(minResult, 5.0);
+
+ Object maxResult =
AggregationFunctionUtils.getAggregationResult(mockFunction(AggregationFunctionType.MAX),
+ dataSource, 100, "TEST");
+ assertEquals(maxResult, 10.0);
+ }
+
+ @Test
+ public void testUnsupportedFunctionThrows() {
Review Comment:
Good start. This only exercises COUNT/MIN/MAX/MODE/null though - the
DISTINCT* branches are the main beneficiaries of the non-scan path and they do
real dictionary iteration (and BYTES -> HLL merging), so a case that resolves
e.g. DISTINCTCOUNT from a real dictionary would add meaningful coverage.
##########
pinot-core/src/main/java/org/apache/pinot/core/plan/AggregationPlanNode.java:
##########
@@ -173,40 +197,68 @@ private boolean hasNullValues(AggregationFunction[]
aggregationFunctions) {
}
/**
- * Returns {@code true} if the given aggregations can be solved with
dictionary or column metadata, {@code false}
- * otherwise.
+ * Returns {@code true} if the given aggregation function can be resolved
from the column dictionary or metadata
+ * (without scanning the segment), {@code false} otherwise. {@code COUNT} is
always eligible. Functions whose result
+ * is derived numerically from the column min/max (e.g. MIN, MAX,
MINMAXRANGE) are only eligible for numeric columns,
+ * since non-numeric columns (e.g. BYTES) store min/max as raw values that
cannot be parsed as numbers.
+ *
+ * @param aggregationFunction aggregation function to test
+ * @param dataSource the function argument's data source (see {@link
#getDataSourceForAggregationFunction})
*/
- private boolean isFitForNonScanBasedPlan() {
- AggregationFunction[] aggregationFunctions =
_queryContext.getAggregationFunctions();
- assert aggregationFunctions != null;
- for (AggregationFunction<?, ?> aggregationFunction : aggregationFunctions)
{
- if (aggregationFunction.getType() == COUNT) {
- continue;
- }
- ExpressionContext argument =
aggregationFunction.getInputExpressions().get(0);
- if (argument.getType() != ExpressionContext.Type.IDENTIFIER) {
- return false;
- }
- DataSource dataSource =
_indexSegment.getDataSource(argument.getIdentifier(),
_queryContext.getSchema());
- if (DICTIONARY_BASED_FUNCTIONS.contains(aggregationFunction.getType())) {
- if (dataSource.getDictionary() != null) {
- continue;
- }
- }
- if (METADATA_BASED_FUNCTIONS.contains(aggregationFunction.getType())) {
- if (dataSource.getDataSourceMetadata().getMaxValue() != null
- && dataSource.getDataSourceMetadata().getMinValue() != null) {
- continue;
- }
- }
+ private boolean isFitForNonScanBasedPlan(AggregationFunction<?, ?>
aggregationFunction,
+ @Nullable DataSource dataSource) {
+ AggregationFunctionType functionType = aggregationFunction.getType();
+ if (functionType == COUNT) {
+ return true;
+ }
+
+ if (dataSource == null) {
+ // Aggregation function does not have a single identifier argument (e.g.
COUNT(*) or COUNT(1)),
+ // so it cannot be resolved from metadata
return false;
}
- return true;
+
+ // MIN/MAX/MINMAXRANGE derive their result numerically from the column
min/max, which is only valid for numeric
+ // columns. Non-numeric columns (e.g. BYTES) store min/max as raw values
that cannot be parsed as numbers.
+ if (NUMERIC_METADATA_FUNCTIONS.contains(functionType)
+ &&
!dataSource.getDataSourceMetadata().getDataType().getStoredType().isNumeric()) {
+ return false;
+ }
+
+ if (dataSource.getDictionary() != null &&
DICTIONARY_BASED_FUNCTIONS.contains(functionType)) {
+ return true;
+ }
+
+ return METADATA_BASED_FUNCTIONS.contains(functionType)
+ && dataSource.getDataSourceMetadata().getMaxValue() != null
+ && dataSource.getDataSourceMetadata().getMinValue() != null;
}
private static boolean canOptimizeFilteredCount(BaseFilterOperator
filterOperator,
AggregationFunction[] aggregationFunctions) {
return (aggregationFunctions.length == 1 &&
aggregationFunctions[0].getType() == COUNT)
&& filterOperator.canOptimizeCount();
}
+
+ /**
+ * Returns the data source for the given aggregation function's argument, or
{@code null} if the function has no
+ * argument (e.g. {@code COUNT(*)}) or its argument is not a single column
identifier (e.g. {@code COUNT(1)} or a
+ * transform expression), in which case it cannot be resolved from
dictionary/metadata.
+ *
+ * @param aggregationFunction aggregation function whose argument data
source is resolved
+ * @return the argument's data source, or {@code null} if it has no single
identifier argument
+ */
+ @Nullable
+ private DataSource
getDataSourceForAggregationFunction(AggregationFunction<?, ?>
aggregationFunction) {
+ List<?> inputExpressions = aggregationFunction.getInputExpressions();
+ if (!inputExpressions.isEmpty()) {
+ ExpressionContext argument =
aggregationFunction.getInputExpressions().get(0);
Review Comment:
nit: can reuse the `inputExpressions` local here instead of calling
`getInputExpressions()` a second time.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
