yashmayya commented on code in PR #19510: URL: https://github.com/apache/pinot/pull/19510#discussion_r4030636619
########## pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotApproximateAggregateRewriteRule.java: ########## @@ -0,0 +1,186 @@ +/** + * 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.calcite.rel.rules; + +import java.util.ArrayList; +import java.util.List; +import java.util.Objects; +import javax.annotation.Nullable; +import org.apache.calcite.plan.Context; +import org.apache.calcite.plan.RelOptRule; +import org.apache.calcite.plan.RelOptRuleCall; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.core.Aggregate; +import org.apache.calcite.rel.core.AggregateCall; +import org.apache.calcite.rel.logical.LogicalAggregate; +import org.apache.calcite.rel.logical.LogicalProject; +import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.sql.SqlAggFunction; +import org.apache.calcite.sql.SqlFunctionCategory; +import org.apache.calcite.sql.SqlKind; +import org.apache.calcite.sql.type.ReturnTypes; +import org.apache.calcite.tools.RelBuilderFactory; +import org.apache.pinot.common.function.sql.PinotSqlAggFunction; +import org.apache.pinot.query.QueryEnvironment; +import org.apache.pinot.query.context.PlannerContext; +import org.apache.pinot.segment.spi.AggregationFunctionType; +import org.apache.pinot.spi.utils.CommonConstants.Broker.PlannerRuleNames; + + +/// Rewrites exact aggregations into their threshold-based approximate counterparts, so that one cluster config can +/// stop unbounded per-group accumulators from taking servers down: +/// - `DISTINCT_COUNT(x)` and `COUNT(DISTINCT x)` -> `DISTINCT_COUNT_SMART_HLL(x)` +/// - `PERCENTILE(x, p)` -> `PERCENTILE_SMART_TDIGEST(x, p)` +/// +/// This is the multi-stage counterpart of `BaseSingleStageBrokerRequestHandler.handleApproximateFunctionOverride`, +/// gated by the same resolved setting, read from [QueryEnvironment.Config#useApproximateFunction()]. The rewrite keeps +/// the original return type, so switching the config on does not change the result schema. +/// +/// It must run before [PinotAggregateExchangeNodeInsertRule], which derives the leaf-to-final intermediate result +/// format from the function name, hence `Phase.BASIC` rather than that rule's `POST_LOGICAL`. That is also why the +/// runtime is not a valid place to do this rewrite. +public class PinotApproximateAggregateRewriteRule extends RelOptRule { + public static final PinotApproximateAggregateRewriteRule INSTANCE = + new PinotApproximateAggregateRewriteRule(PinotRuleUtils.PINOT_REL_FACTORY); + + private PinotApproximateAggregateRewriteRule(RelBuilderFactory factory) { + super(operand(LogicalAggregate.class, any()), factory, PlannerRuleNames.APPROXIMATE_AGGREGATE_REWRITE); + } + + @Override + public boolean matches(RelOptRuleCall call) { + QueryEnvironment.Config envConfig = envConfig(call); + if (envConfig == null || !envConfig.useApproximateFunction()) { + return false; + } + // Requiring a rewritable call is also what makes the rule terminate: the rewritten names no longer match. + Aggregate aggRel = call.rel(0); + return aggRel.getAggCallList().stream().anyMatch(aggCall -> targetOf(aggCall) != null); + } + + @Override + public void onMatch(RelOptRuleCall call) { + Aggregate aggRel = call.rel(0); + QueryEnvironment.Config envConfig = Objects.requireNonNull(envConfig(call)); + RelNode input = aggRel.getInput(); + int numInputFields = input.getRowType().getFieldCount(); + + boolean hasDistinctCount = false; + boolean hasPercentile = false; + for (AggregateCall aggCall : aggRel.getAggCallList()) { + AggregationFunctionType target = targetOf(aggCall); + hasDistinctCount |= target == AggregationFunctionType.DISTINCTCOUNTSMARTHLL; + hasPercentile |= target == AggregationFunctionType.PERCENTILESMARTTDIGEST; + } + String distinctCountParams = hasDistinctCount ? envConfig.approximateFunctionDistinctCountParams() : ""; + String percentileParams = hasPercentile ? envConfig.approximateFunctionPercentileParams() : ""; + + // The parameters must be input fields, because an aggregate call holds field indices rather than inline literals, + // so they are projected underneath the aggregate. Placing them in the immediate project is what lets + // PinotAggregateExchangeNodeInsertRule inline them back into the pushed-down call. With no parameters configured + // the plan keeps the shape it had before, with no extra project. + RelNode newInput = input; + int distinctCountParamsIndex = -1; + int percentileParamsIndex = -1; + if (!distinctCountParams.isEmpty() || !percentileParams.isEmpty()) { + RexBuilder rexBuilder = aggRel.getCluster().getRexBuilder(); + List<RexNode> projects = new ArrayList<>(numInputFields + 2); + for (int i = 0; i < numInputFields; i++) { + projects.add(rexBuilder.makeInputRef(input, i)); + } + if (!distinctCountParams.isEmpty()) { + distinctCountParamsIndex = projects.size(); + projects.add(rexBuilder.makeLiteral(distinctCountParams)); + } + if (!percentileParams.isEmpty()) { + percentileParamsIndex = projects.size(); + projects.add(rexBuilder.makeLiteral(percentileParams)); + } + newInput = LogicalProject.create(input, List.of(), projects, (List<String>) null); + } + + List<AggregateCall> rewrittenAggCalls = new ArrayList<>(aggRel.getAggCallList().size()); + for (AggregateCall aggCall : aggRel.getAggCallList()) { + AggregationFunctionType target = targetOf(aggCall); + if (target == null) { + rewrittenAggCalls.add(aggCall); + continue; + } + int paramsIndex = target == AggregationFunctionType.DISTINCTCOUNTSMARTHLL + ? distinctCountParamsIndex : percentileParamsIndex; + List<Integer> argList = aggCall.getArgList(); + if (paramsIndex >= 0) { + argList = new ArrayList<>(argList); + argList.add(paramsIndex); + } + rewrittenAggCalls.add(rewrite(aggCall, target, argList, newInput, aggRel.getGroupCount())); + } + + PlannerContext plannerContext = call.getPlanner().getContext().unwrap(PlannerContext.class); + if (plannerContext != null) { + plannerContext.setApproximateFunctionApplied(); + } + call.transformTo( + aggRel.copy(aggRel.getTraitSet(), newInput, aggRel.getGroupSet(), aggRel.getGroupSets(), rewrittenAggCalls)); + } + + /// Returns the approximate function this call should become, or `null` when the call must be left alone. + @Nullable + private static AggregationFunctionType targetOf(AggregateCall aggCall) { + SqlAggFunction aggFunction = aggCall.getAggregation(); + if (aggCall.isDistinct()) { + // COUNT(DISTINCT x) is the standard SQL spelling of DISTINCT_COUNT, and PinotAggregateExchangeNodeInsertRule + // only renames it in POST_LOGICAL, after this rule, so it has to be matched on the kind here. Multi-argument + // COUNT(DISTINCT a, b) means something else and is left alone. + return aggFunction.getKind() == SqlKind.COUNT && aggCall.getArgList().size() == 1 + ? AggregationFunctionType.DISTINCTCOUNTSMARTHLL : null; + } + // The smart functions take multi-valued input themselves, so the MV spellings map onto the same targets, which + // keeps this in step with the single-stage rewrite. + String name = AggregationFunctionType.getNormalizedAggregationFunctionName(aggFunction.getName()); + if (name.equals(AggregationFunctionType.DISTINCTCOUNT.name()) + || name.equals(AggregationFunctionType.DISTINCTCOUNTMV.name())) { + return AggregationFunctionType.DISTINCTCOUNTSMARTHLL; + } + if (name.equals(AggregationFunctionType.PERCENTILE.name()) + || name.equals(AggregationFunctionType.PERCENTILEMV.name())) { + return AggregationFunctionType.PERCENTILESMARTTDIGEST; + } + return null; + } + + private static AggregateCall rewrite(AggregateCall aggCall, AggregationFunctionType target, List<Integer> argList, + RelNode input, int groupCount) { + // Pinning the original type keeps the rewrite invisible to the caller: PERCENTILE infers ARG0 while + // PERCENTILE_SMART_TDIGEST infers DOUBLE, and a silent cluster-wide rewrite must not change the result schema. + SqlAggFunction newAggFunction = new PinotSqlAggFunction(target.name(), SqlKind.OTHER_FUNCTION, Review Comment: Good catch, that is a real truncation. Declared \`SqlTypeName.DOUBLE\` as the final return type on \`PERCENTILESMARTTDIGEST\`, with a comment saying it is there because of the pinning. Agreed the fallback fix belongs in its own diff. ########## pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotApproximateAggregateRewriteRule.java: ########## @@ -0,0 +1,186 @@ +/** + * 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.calcite.rel.rules; + +import java.util.ArrayList; +import java.util.List; +import java.util.Objects; +import javax.annotation.Nullable; +import org.apache.calcite.plan.Context; +import org.apache.calcite.plan.RelOptRule; +import org.apache.calcite.plan.RelOptRuleCall; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.core.Aggregate; +import org.apache.calcite.rel.core.AggregateCall; +import org.apache.calcite.rel.logical.LogicalAggregate; +import org.apache.calcite.rel.logical.LogicalProject; +import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.sql.SqlAggFunction; +import org.apache.calcite.sql.SqlFunctionCategory; +import org.apache.calcite.sql.SqlKind; +import org.apache.calcite.sql.type.ReturnTypes; +import org.apache.calcite.tools.RelBuilderFactory; +import org.apache.pinot.common.function.sql.PinotSqlAggFunction; +import org.apache.pinot.query.QueryEnvironment; +import org.apache.pinot.query.context.PlannerContext; +import org.apache.pinot.segment.spi.AggregationFunctionType; +import org.apache.pinot.spi.utils.CommonConstants.Broker.PlannerRuleNames; + + +/// Rewrites exact aggregations into their threshold-based approximate counterparts, so that one cluster config can +/// stop unbounded per-group accumulators from taking servers down: +/// - `DISTINCT_COUNT(x)` and `COUNT(DISTINCT x)` -> `DISTINCT_COUNT_SMART_HLL(x)` +/// - `PERCENTILE(x, p)` -> `PERCENTILE_SMART_TDIGEST(x, p)` +/// +/// This is the multi-stage counterpart of `BaseSingleStageBrokerRequestHandler.handleApproximateFunctionOverride`, +/// gated by the same resolved setting, read from [QueryEnvironment.Config#useApproximateFunction()]. The rewrite keeps +/// the original return type, so switching the config on does not change the result schema. +/// +/// It must run before [PinotAggregateExchangeNodeInsertRule], which derives the leaf-to-final intermediate result +/// format from the function name, hence `Phase.BASIC` rather than that rule's `POST_LOGICAL`. That is also why the +/// runtime is not a valid place to do this rewrite. +public class PinotApproximateAggregateRewriteRule extends RelOptRule { + public static final PinotApproximateAggregateRewriteRule INSTANCE = + new PinotApproximateAggregateRewriteRule(PinotRuleUtils.PINOT_REL_FACTORY); + + private PinotApproximateAggregateRewriteRule(RelBuilderFactory factory) { + super(operand(LogicalAggregate.class, any()), factory, PlannerRuleNames.APPROXIMATE_AGGREGATE_REWRITE); + } + + @Override + public boolean matches(RelOptRuleCall call) { + QueryEnvironment.Config envConfig = envConfig(call); + if (envConfig == null || !envConfig.useApproximateFunction()) { + return false; + } + // Requiring a rewritable call is also what makes the rule terminate: the rewritten names no longer match. + Aggregate aggRel = call.rel(0); + return aggRel.getAggCallList().stream().anyMatch(aggCall -> targetOf(aggCall) != null); + } + + @Override + public void onMatch(RelOptRuleCall call) { + Aggregate aggRel = call.rel(0); + QueryEnvironment.Config envConfig = Objects.requireNonNull(envConfig(call)); + RelNode input = aggRel.getInput(); + int numInputFields = input.getRowType().getFieldCount(); + + boolean hasDistinctCount = false; + boolean hasPercentile = false; + for (AggregateCall aggCall : aggRel.getAggCallList()) { + AggregationFunctionType target = targetOf(aggCall); + hasDistinctCount |= target == AggregationFunctionType.DISTINCTCOUNTSMARTHLL; + hasPercentile |= target == AggregationFunctionType.PERCENTILESMARTTDIGEST; + } + String distinctCountParams = hasDistinctCount ? envConfig.approximateFunctionDistinctCountParams() : ""; + String percentileParams = hasPercentile ? envConfig.approximateFunctionPercentileParams() : ""; + + // The parameters must be input fields, because an aggregate call holds field indices rather than inline literals, + // so they are projected underneath the aggregate. Placing them in the immediate project is what lets + // PinotAggregateExchangeNodeInsertRule inline them back into the pushed-down call. With no parameters configured + // the plan keeps the shape it had before, with no extra project. + RelNode newInput = input; + int distinctCountParamsIndex = -1; + int percentileParamsIndex = -1; + if (!distinctCountParams.isEmpty() || !percentileParams.isEmpty()) { + RexBuilder rexBuilder = aggRel.getCluster().getRexBuilder(); + List<RexNode> projects = new ArrayList<>(numInputFields + 2); + for (int i = 0; i < numInputFields; i++) { + projects.add(rexBuilder.makeInputRef(input, i)); + } + if (!distinctCountParams.isEmpty()) { + distinctCountParamsIndex = projects.size(); + projects.add(rexBuilder.makeLiteral(distinctCountParams)); + } + if (!percentileParams.isEmpty()) { + percentileParamsIndex = projects.size(); + projects.add(rexBuilder.makeLiteral(percentileParams)); + } + newInput = LogicalProject.create(input, List.of(), projects, (List<String>) null); Review Comment: Added a comment naming the invariant and both consumers. Also added a plan test over a join and a subquery — both hold. ########## pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/BaseDistinctCountSmartSketchAggregationFunction.java: ########## @@ -383,6 +387,7 @@ public final void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByRes throw getIllegalDataTypeException(valueType, false); } } + checkAndConvertValueSetsForGroups(groupByResultHolder, groupKeyArray, length, storedType); Review Comment: Fair, that cost landed on everyone. Moved the check into \`getValueSet\`, which already fetched the set, so the second pass is gone and only non-null rows are visited. It now converts on the first touch after the threshold is crossed, which bounds a group at threshold + one batch. ########## pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/BaseDistinctCountSmartSketchAggregationFunction.java: ########## @@ -61,6 +61,10 @@ protected BaseDistinctCountSmartSketchAggregationFunction(ExpressionContext expr protected abstract Object convertSetToSketch(Set valueSet, DataType storedType); + /// Adds every value of the set into a sketch that [#convertSetToSketch] already created. Used to fold the values + /// that arrive after a group has been converted, without building a second sketch. + protected abstract void addSetToSketch(Object sketch, Set valueSet, DataType storedType); Review Comment: Added a test over all three sketch families and STRING columns, with counts large enough that a dropped batch would halve the estimate. ULL needed that — it estimates 5 for 6 distinct values. ########## pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java: ########## @@ -500,10 +500,32 @@ public static class Broker { public static final String DISABLE_GROOVY = "pinot.broker.disable.query.groovy"; public static final boolean DEFAULT_DISABLE_GROOVY = true; - // Rewrite potential expensive functions to their approximation counterparts - // - DISTINCT_COUNT -> DISTINCT_COUNT_SMART_HLL - // - PERCENTILE -> PERCENTILE_SMART_TDIGEST + /// Rewrite potential expensive functions to their approximation counterparts, in both query engines: + /// - DISTINCT_COUNT and COUNT(DISTINCT) -> DISTINCT_COUNT_SMART_HLL + /// - PERCENTILE -> PERCENTILE_SMART_TDIGEST + /// + /// The rewritten functions stay exact until an accumulator exceeds their conversion threshold, so this bounds + /// server memory without changing the answer for low-cardinality inputs. + /// + /// Settable in the broker conf, read once at startup, or in the Helix cluster config, which wins and which + /// brokers pick up without a restart. Both are defaults: the `useApproximateFunction` query option overrides + /// them, and so does `QueryConfig.useApproximateFunction`, though only in the single-stage engine, because a + /// multi-stage query can span tables and so resolves the setting before it knows the table set. public static final String USE_APPROXIMATE_FUNCTION = "pinot.broker.use.approximate.function"; + public static final boolean DEFAULT_USE_APPROXIMATE_FUNCTION = false; + + /// Parameters passed verbatim as the trailing argument of the calls the rewrite produces, for example + /// `threshold=10000;log2m=12;dictThreshold=10000` and `threshold=1000;compression=100`. Empty means no argument + /// is added, so the aggregation function defaults apply. The two functions reject each other's parameter names, + /// hence one key each. + /// + /// Worst-case memory for a group-by is `threshold` values per group, so size this together with + /// `pinot.server.query.executor.num.groups.limit` rather than in isolation. + public static final String APPROXIMATE_FUNCTION_DISTINCT_COUNT_PARAMS = Review Comment: Right, I only documented one direction. Reworded to say a low threshold trades raw values for eagerly allocated registers and can raise memory, and to size it with \`log2m\`/\`compression\`. ########## pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/ApproximateFunctionOverrideProvider.java: ########## @@ -0,0 +1,145 @@ +/** + * 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.broker.requesthandler; + +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.function.Consumer; +import javax.annotation.Nullable; +import org.apache.pinot.common.request.Literal; +import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.core.query.aggregation.function.DistinctCountSmartHLLAggregationFunction; +import org.apache.pinot.core.query.aggregation.function.PercentileSmartTDigestAggregationFunction; +import org.apache.pinot.spi.config.provider.PinotClusterConfigChangeListener; +import org.apache.pinot.spi.env.PinotConfiguration; +import org.apache.pinot.spi.utils.CommonConstants.Broker; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +/// Resolves the defaults for the approximate function rewrite, from the broker conf and from the Helix cluster config. +/// The cluster config wins where it sets a key, and it is applied live: registering this as a +/// [PinotClusterConfigChangeListener] means a change takes effect on the next query without a broker restart. +/// +/// Thread-safe. [#getSettings()] hands out one immutable snapshot, so a query that reads it once cannot see a config +/// change half applied. +public class ApproximateFunctionOverrideProvider implements PinotClusterConfigChangeListener { + private static final Logger LOGGER = LoggerFactory.getLogger(ApproximateFunctionOverrideProvider.class); + + private static final ExpressionContext PROBE_COLUMN = ExpressionContext.forIdentifier("__probe__"); + private static final ExpressionContext PROBE_PERCENTILE = ExpressionContext.forLiteral(Literal.doubleValue(50.0)); + + private final Settings _brokerConfSettings; + private volatile Settings _settings; + + public ApproximateFunctionOverrideProvider(PinotConfiguration config) { + _brokerConfSettings = new Settings( + config.getProperty(Broker.USE_APPROXIMATE_FUNCTION, Broker.DEFAULT_USE_APPROXIMATE_FUNCTION), + validated(config.getProperty(Broker.APPROXIMATE_FUNCTION_DISTINCT_COUNT_PARAMS, + Broker.DEFAULT_APPROXIMATE_FUNCTION_PARAMS), Broker.APPROXIMATE_FUNCTION_DISTINCT_COUNT_PARAMS, + Broker.DEFAULT_APPROXIMATE_FUNCTION_PARAMS, ApproximateFunctionOverrideProvider::probeDistinctCount), + validated(config.getProperty(Broker.APPROXIMATE_FUNCTION_PERCENTILE_PARAMS, + Broker.DEFAULT_APPROXIMATE_FUNCTION_PARAMS), Broker.APPROXIMATE_FUNCTION_PERCENTILE_PARAMS, + Broker.DEFAULT_APPROXIMATE_FUNCTION_PARAMS, ApproximateFunctionOverrideProvider::probePercentile)); + _settings = _brokerConfSettings; + } + + @Override + public void onChange(Set<String> changedConfigs, Map<String, String> clusterConfigs) { + if (!changedConfigs.contains(Broker.USE_APPROXIMATE_FUNCTION) + && !changedConfigs.contains(Broker.APPROXIMATE_FUNCTION_DISTINCT_COUNT_PARAMS) + && !changedConfigs.contains(Broker.APPROXIMATE_FUNCTION_PERCENTILE_PARAMS)) { + return; + } + // A key the cluster config does not set, or that was removed from it, falls back to the broker conf value. The + // fallback is the broker conf rather than the last good live value, so a broker that restarts resolves the same. + String enabled = clusterConfigs.get(Broker.USE_APPROXIMATE_FUNCTION); Review Comment: Done. Anything that is not true/false now logs and keeps the broker conf value. ########## pinot-common/src/main/java/org/apache/pinot/common/response/BrokerResponse.java: ########## @@ -311,6 +311,15 @@ default long getRealtimeTotalMemAllocatedBytes() { /// @return true if RLS filters were applied, false otherwise boolean getRLSFiltersApplied(); + /// Set whether the broker rewrote an exact aggregation into its approximate counterpart, for example + /// `DISTINCT_COUNT` into `DISTINCT_COUNT_SMART_HLL`. The result is then approximate, not exact. + /// @param approximateFunctionApplied true if at least one function was rewritten + void setApproximateFunctionApplied(boolean approximateFunctionApplied); Review Comment: Done, both are \`default\` now. -- 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]
