Skip to content

Improve input validation for approx_percentile #23546

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
losipiuk merged 2 commits into from
Sep 25, 2024
Merged

Improve input validation for approx_percentile #23546

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
73e85f5
Ensure no unnecessary boxing happens
losipiuk Sep 24, 2024
152fef6
Improve input validation for approx_percentile
losipiuk Sep 24, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions .../src/main/java/io/trino/operator/aggregation/ApproximateDoublePercentileAggregations.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
import io.trino.spi.type.StandardTypes;

import static com.google.common.base.Preconditions.checkState;
import static io.trino.operator.scalar.TDigestFunctions.verifyValue;
import static io.trino.operator.scalar.TDigestFunctions.verifyWeight;
import static io.trino.spi.StandardErrorCode.INVALID_FUNCTION_ARGUMENT;
import static io.trino.spi.type.DoubleType.DOUBLE;
Expand All @@ -38,6 +39,7 @@ private ApproximateDoublePercentileAggregations() {}
@InputFunction
public static void input(@AggregationState TDigestAndPercentileState state, @SqlType(StandardTypes.DOUBLE) double value, @SqlType(StandardTypes.DOUBLE) double percentile)
{
verifyValue(value);
TDigest digest = state.getDigest();

if (digest == null) {
Expand All @@ -57,6 +59,7 @@ public static void input(@AggregationState TDigestAndPercentileState state, @Sql
@InputFunction
public static void weightedInput(@AggregationState TDigestAndPercentileState state, @SqlType(StandardTypes.DOUBLE) double value, @SqlType(StandardTypes.DOUBLE) double weight, @SqlType(StandardTypes.DOUBLE) double percentile)
{
verifyValue(value);
verifyWeight(weight);

TDigest digest = state.getDigest();
Expand Down
4 changes: 4 additions & 0 deletions ...main/java/io/trino/operator/aggregation/ApproximateDoublePercentileArrayAggregations.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@

import java.util.List;

import static io.trino.operator.scalar.TDigestFunctions.verifyValue;
import static io.trino.operator.scalar.TDigestFunctions.verifyWeight;
import static io.trino.spi.StandardErrorCode.INVALID_FUNCTION_ARGUMENT;
import static io.trino.spi.type.DoubleType.DOUBLE;
Expand All @@ -43,6 +44,8 @@ private ApproximateDoublePercentileArrayAggregations() {}
@InputFunction
public static void input(@AggregationState TDigestAndPercentileArrayState state, @SqlType(StandardTypes.DOUBLE) double value, @SqlType("array(double)") Block percentilesArrayBlock)
{
verifyValue(value);

initializePercentilesArray(state, percentilesArrayBlock);
initializeDigest(state);

Expand All @@ -55,6 +58,7 @@ public static void input(@AggregationState TDigestAndPercentileArrayState state,
@InputFunction
public static void weightedInput(@AggregationState TDigestAndPercentileArrayState state, @SqlType(StandardTypes.DOUBLE) double value, @SqlType(StandardTypes.DOUBLE) double weight, @SqlType("array(double)") Block percentilesArrayBlock)
{
verifyValue(value);
verifyWeight(weight);

initializePercentilesArray(state, percentilesArrayBlock);
Expand Down
2 changes: 1 addition & 1 deletion ...in/src/main/java/io/trino/operator/aggregation/ApproximateLongPercentileAggregations.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,7 @@ public static void output(@AggregationState TDigestAndPercentileState state, Blo
public static double toDoubleExact(long value)
{
double doubleValue = (double) value;
checkCondition((long) doubleValue == value, INVALID_FUNCTION_ARGUMENT, "no exact double representation for long: %s", value);
checkCondition((long) doubleValue == value, INVALID_FUNCTION_ARGUMENT, () -> String.format("no exact double representation for long: %s", value));
return doubleValue;
}
}
3 changes: 2 additions & 1 deletion .../main/java/io/trino/operator/aggregation/LegacyApproximateLongPercentileAggregations.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@

import io.airlift.stats.QuantileDigest;
import io.trino.operator.aggregation.state.QuantileDigestAndPercentileState;
import io.trino.spi.TrinoException;
import io.trino.spi.block.BlockBuilder;
import io.trino.spi.function.AggregationFunction;
import io.trino.spi.function.AggregationState;
Expand Down Expand Up @@ -50,7 +51,7 @@ public static void weightedInput(@AggregationState QuantileDigestAndPercentileSt
digest = new QuantileDigest(accuracy);
}
else {
throw new IllegalArgumentException("Percentile accuracy must be strictly between 0 and 1");
throw new TrinoException(INVALID_FUNCTION_ARGUMENT, "Percentile accuracy must be strictly between 0 and 1");
}
state.setDigest(digest);
state.addMemoryUsage(digest.estimatedInMemorySizeInBytes());
Expand Down
4 changes: 2 additions & 2 deletions core/trino-main/src/main/java/io/trino/operator/scalar/ArrayCombinationsFunction.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,8 +56,8 @@ public static Block combinations(
{
int arrayLength = array.getPositionCount();
int combinationLength = toIntExact(n);
checkCondition(combinationLength >= 0, INVALID_FUNCTION_ARGUMENT, "combination size must not be negative: %s", combinationLength);
checkCondition(combinationLength <= MAX_COMBINATION_LENGTH, INVALID_FUNCTION_ARGUMENT, "combination size must not exceed %s: %s", MAX_COMBINATION_LENGTH, combinationLength);
checkCondition(combinationLength >= 0, INVALID_FUNCTION_ARGUMENT, () -> String.format("combination size must not be negative: %s", combinationLength));
checkCondition(combinationLength <= MAX_COMBINATION_LENGTH, INVALID_FUNCTION_ARGUMENT, () -> String.format("combination size must not exceed %s: %s", MAX_COMBINATION_LENGTH, combinationLength));

ArrayType arrayType = new ArrayType(elementType);
if (combinationLength > arrayLength) {
Expand Down
4 changes: 2 additions & 2 deletions core/trino-main/src/main/java/io/trino/operator/scalar/ArrayTrimFunction.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,8 +38,8 @@ public static Block trim(
@SqlType("array(E)") Block array,
@SqlType(StandardTypes.BIGINT) long size)
{
checkCondition(size >= 0, INVALID_FUNCTION_ARGUMENT, "size must not be negative: %s", size);
checkCondition(size <= array.getPositionCount(), INVALID_FUNCTION_ARGUMENT, "size must not exceed array cardinality %s: %s", array.getPositionCount(), size);
checkCondition(size >= 0, INVALID_FUNCTION_ARGUMENT, () -> String.format("size must not be negative: %s", size));
checkCondition(size <= array.getPositionCount(), INVALID_FUNCTION_ARGUMENT, () -> String.format("size must not exceed array cardinality %s: %s", array.getPositionCount(), size));

return array.getRegion(0, toIntExact(array.getPositionCount() - size));
}
Expand Down
4 changes: 2 additions & 2 deletions core/trino-main/src/main/java/io/trino/operator/scalar/QuantileDigestFunctions.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -141,13 +141,13 @@ public static Block valuesAtQuantilesBigint(@SqlType("qdigest(bigint)") Slice in

public static double verifyAccuracy(double accuracy)
{
checkCondition(accuracy > 0 && accuracy < 1, INVALID_FUNCTION_ARGUMENT, "Percentile accuracy must be exclusively between 0 and 1, was %s", accuracy);
checkCondition(accuracy > 0 && accuracy < 1, INVALID_FUNCTION_ARGUMENT, () -> String.format("Percentile accuracy must be exclusively between 0 and 1, was %s", accuracy));
return accuracy;
}

public static long verifyWeight(long weight)
{
checkCondition(weight > 0, INVALID_FUNCTION_ARGUMENT, "Percentile weight must be > 0, was %s", weight);
checkCondition(weight > 0, INVALID_FUNCTION_ARGUMENT, () -> String.format("Percentile weight must be > 0, was %s", weight));
return weight;
}
}
8 changes: 7 additions & 1 deletion core/trino-main/src/main/java/io/trino/operator/scalar/TDigestFunctions.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,9 +59,15 @@ public static Block valuesAtQuantiles(@SqlType(StandardTypes.TDIGEST) TDigest in
return output.build();
}

public static void verifyValue(double value)
{
checkCondition(Double.isFinite(value), INVALID_FUNCTION_ARGUMENT, () -> String.format("value must be finite; was %s", value));
}

public static double verifyWeight(double weight)
{
checkCondition(weight >= 1, INVALID_FUNCTION_ARGUMENT, "weight must be >= 1, was %s", weight);
checkCondition(Double.isFinite(weight), INVALID_FUNCTION_ARGUMENT, () -> String.format("weight must be finite, was %s", weight));
checkCondition(weight >= 1, INVALID_FUNCTION_ARGUMENT, () -> String.format("weight must be >= 1, was %s", weight));
return weight;
}
}
4 changes: 2 additions & 2 deletions core/trino-main/src/main/java/io/trino/type/DecimalCasts.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -573,7 +573,7 @@ public static Int128 jsonToLongDecimal(Slice json, long precision, long scale, I
try (JsonParser parser = createJsonParser(JSON_FACTORY, json)) {
parser.nextToken();
Int128 result = currentTokenAsLongDecimal(parser, intPrecision(precision), DecimalConversions.intScale(scale));
checkCondition(parser.nextToken() == null, INVALID_CAST_ARGUMENT, "Cannot cast input json to DECIMAL(%s,%s)", precision, scale); // check no trailing token
checkCondition(parser.nextToken() == null, INVALID_CAST_ARGUMENT, () -> String.format("Cannot cast input json to DECIMAL(%s,%s)", precision, scale)); // check no trailing token
return result;
}
catch (IOException | NumberFormatException | JsonCastException e) {
Expand All @@ -587,7 +587,7 @@ public static Long jsonToShortDecimal(Slice json, long precision, long scale, lo
try (JsonParser parser = createJsonParser(JSON_FACTORY, json)) {
parser.nextToken();
Long result = currentTokenAsShortDecimal(parser, intPrecision(precision), DecimalConversions.intScale(scale));
checkCondition(parser.nextToken() == null, INVALID_CAST_ARGUMENT, "Cannot cast input json to DECIMAL(%s,%s)", precision, scale); // check no trailing token
checkCondition(parser.nextToken() == null, INVALID_CAST_ARGUMENT, () -> String.format("Cannot cast input json to DECIMAL(%s,%s)", precision, scale)); // check no trailing token
return result;
}
catch (IOException | NumberFormatException | JsonCastException e) {
Expand Down
8 changes: 7 additions & 1 deletion core/trino-main/src/main/java/io/trino/util/Failures.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
*/
package io.trino.util;

import com.google.common.base.Supplier;
import com.google.common.collect.ImmutableList;
import com.google.errorprone.annotations.FormatMethod;
import io.trino.client.ErrorLocation;
Expand Down Expand Up @@ -60,9 +61,14 @@ public static ExecutionFailureInfo toFailure(Throwable failure)

@FormatMethod
public static void checkCondition(boolean condition, ErrorCodeSupplier errorCode, String formatString, Object... args)
{
checkCondition(condition, errorCode, () -> format(formatString, args));
}

public static void checkCondition(boolean condition, ErrorCodeSupplier errorCode, Supplier<String> errorMessage)
{
if (!condition) {
throw new TrinoException(errorCode, format(formatString, args));
throw new TrinoException(errorCode, errorMessage.get());
}
}

Expand Down
7 changes: 7 additions & 0 deletions core/trino-main/src/test/java/io/trino/operator/aggregation/AggregationTestUtils.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
import io.trino.spi.type.BooleanType;
import io.trino.spi.type.Type;
import io.trino.sql.analyzer.TypeSignatureProvider;
import io.trino.testing.assertions.TrinoExceptionAssert;
import org.apache.commons.math3.util.Precision;

import java.util.Arrays;
Expand All @@ -38,6 +39,7 @@
import static io.trino.sql.planner.plan.AggregationNode.Step.FINAL;
import static io.trino.sql.planner.plan.AggregationNode.Step.PARTIAL;
import static io.trino.sql.planner.plan.AggregationNode.Step.SINGLE;
import static io.trino.testing.assertions.TrinoExceptionAssert.assertTrinoExceptionThrownBy;
import static java.lang.String.format;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Fail.fail;
Expand All @@ -58,6 +60,11 @@ public static void assertAggregation(TestingFunctionResolution functionResolutio
assertAggregation(functionResolution, name, parameterTypes, equalAssertion, null, page, expectedValue);
}

public static TrinoExceptionAssert assertAggregationFails(TestingFunctionResolution functionResolution, String name, List<TypeSignatureProvider> parameterTypes, Block... blocks)
{
return assertTrinoExceptionThrownBy(() -> assertAggregation(functionResolution, name, parameterTypes, null, blocks));
}

public static BiFunction<Object, Object, Boolean> makeValidityAssertion(Object expectedValue)
{
if (expectedValue instanceof Double && !expectedValue.equals(Double.NaN)) {
Expand Down
135 changes: 135 additions & 0 deletions ...ain/src/test/java/io/trino/operator/aggregation/TestApproximatePercentileAggregation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@

import com.google.common.collect.ImmutableList;
import io.trino.metadata.TestingFunctionResolution;
import io.trino.spi.TrinoException;
import io.trino.spi.block.ArrayBlockBuilder;
import io.trino.spi.block.Block;
import io.trino.spi.block.BlockBuilder;
Expand All @@ -33,6 +34,8 @@
import static io.trino.block.BlockAssertions.createLongsBlock;
import static io.trino.block.BlockAssertions.createSequenceBlockOfReal;
import static io.trino.operator.aggregation.AggregationTestUtils.assertAggregation;
import static io.trino.operator.aggregation.AggregationTestUtils.assertAggregationFails;
import static io.trino.spi.StandardErrorCode.INVALID_FUNCTION_ARGUMENT;
import static io.trino.spi.type.BigintType.BIGINT;
import static io.trino.spi.type.DoubleType.DOUBLE;
import static io.trino.spi.type.RealType.REAL;
Expand Down Expand Up @@ -442,6 +445,72 @@ public void testFloatPartialStep()
createBlockOfReals(1.0f, 2.0f, 3.0f),
createDoublesBlock(4.0, 2.0, 1.0),
createRleBlock(ImmutableList.of(0.5, 0.8), 3));

// invalid inputs
for (Float invalidValue : List.of(Float.NaN, Float.NEGATIVE_INFINITY, Float.POSITIVE_INFINITY)) {
assertAggregationFails(FUNCTION_RESOLUTION,
"approx_percentile",
FLOAT_APPROXIMATE_PERCENTILE,
createBlockOfReals(invalidValue),
createRleBlock(0.5, 1))
.isInstanceOf(TrinoException.class)
.hasErrorCode(INVALID_FUNCTION_ARGUMENT);

assertAggregationFails(FUNCTION_RESOLUTION,
"approx_percentile",
FLOAT_APPROXIMATE_PERCENTILE_ARRAY,
createBlockOfReals(invalidValue),
createRleBlock(ImmutableList.of(0.5, 0.75), 1))
.isInstanceOf(TrinoException.class)
.hasErrorCode(INVALID_FUNCTION_ARGUMENT);

assertAggregationFails(FUNCTION_RESOLUTION,
"approx_percentile",
FLOAT_APPROXIMATE_PERCENTILE_WEIGHTED,
createBlockOfReals(1.0f),
createDoublesBlock((double) invalidValue),
createDoublesBlock(0.5))
.isInstanceOf(TrinoException.class)
.hasErrorCode(INVALID_FUNCTION_ARGUMENT);

assertAggregationFails(FUNCTION_RESOLUTION,
"approx_percentile",
FLOAT_APPROXIMATE_PERCENTILE_WEIGHTED,
createBlockOfReals(invalidValue),
createDoublesBlock(1.0),
createRleBlock(0.5, 1))
.isInstanceOf(TrinoException.class)
.hasErrorCode(INVALID_FUNCTION_ARGUMENT);

assertAggregationFails(FUNCTION_RESOLUTION,
"approx_percentile",
FLOAT_APPROXIMATE_PERCENTILE_ARRAY_WEIGHTED,
createBlockOfReals(invalidValue),
createDoublesBlock(1.0),
createRleBlock(ImmutableList.of(0.5, 0.75), 1))
.isInstanceOf(TrinoException.class)
.hasErrorCode(INVALID_FUNCTION_ARGUMENT);

assertAggregationFails(FUNCTION_RESOLUTION,
"approx_percentile",
FLOAT_APPROXIMATE_PERCENTILE_ARRAY_WEIGHTED,
createBlockOfReals(1.0f),
createDoublesBlock((double) invalidValue),
createRleBlock(ImmutableList.of(0.5, 0.75), 1))
.isInstanceOf(TrinoException.class)
.hasErrorCode(INVALID_FUNCTION_ARGUMENT);

// for deprecated approx_percentile with accuracy we only validate accuracy for backward compatibility
assertAggregationFails(
FUNCTION_RESOLUTION,
"approx_percentile",
FLOAT_APPROXIMATE_PERCENTILE_WEIGHTED_WITH_ACCURACY,
createBlockOfReals(1.0f),
createDoublesBlock(1.0),
createDoublesBlock(0.99),
createDoublesBlock((double) invalidValue))
.hasErrorCode(INVALID_FUNCTION_ARGUMENT);
}
}

@Test
Expand Down Expand Up @@ -620,6 +689,72 @@ public void testDoublePartialStep()
createDoublesBlock(1.0, 2.0, 3.0),
createDoublesBlock(4.0, 2.0, 1.0),
createRleBlock(ImmutableList.of(0.5, 0.8), 3));

// invalid inputs
for (Double invalidValue : List.of(Double.NaN, Double.NEGATIVE_INFINITY, Double.POSITIVE_INFINITY)) {
assertAggregationFails(FUNCTION_RESOLUTION,
"approx_percentile",
DOUBLE_APPROXIMATE_PERCENTILE,
createDoublesBlock(invalidValue),
createDoublesBlock(0.5))
.isInstanceOf(TrinoException.class)
.hasErrorCode(INVALID_FUNCTION_ARGUMENT);

assertAggregationFails(FUNCTION_RESOLUTION,
"approx_percentile",
DOUBLE_APPROXIMATE_PERCENTILE_ARRAY,
createDoublesBlock(invalidValue),
createRleBlock(ImmutableList.of(0.5, 0.75), 1))
.isInstanceOf(TrinoException.class)
.hasErrorCode(INVALID_FUNCTION_ARGUMENT);

assertAggregationFails(FUNCTION_RESOLUTION,
"approx_percentile",
DOUBLE_APPROXIMATE_PERCENTILE_WEIGHTED,
createDoublesBlock(invalidValue),
createDoublesBlock(1.0),
createRleBlock(0.5, 1))
.isInstanceOf(TrinoException.class)
.hasErrorCode(INVALID_FUNCTION_ARGUMENT);

assertAggregationFails(FUNCTION_RESOLUTION,
"approx_percentile",
DOUBLE_APPROXIMATE_PERCENTILE_WEIGHTED,
createDoublesBlock(1.0),
createDoublesBlock(invalidValue),
createRleBlock(0.5, 1))
.isInstanceOf(TrinoException.class)
.hasErrorCode(INVALID_FUNCTION_ARGUMENT);

assertAggregationFails(FUNCTION_RESOLUTION,
"approx_percentile",
DOUBLE_APPROXIMATE_PERCENTILE_ARRAY_WEIGHTED,
createDoublesBlock(1.0),
createDoublesBlock(invalidValue),
createRleBlock(ImmutableList.of(0.5, 0.75), 1))
.isInstanceOf(TrinoException.class)
.hasErrorCode(INVALID_FUNCTION_ARGUMENT);

assertAggregationFails(FUNCTION_RESOLUTION,
"approx_percentile",
DOUBLE_APPROXIMATE_PERCENTILE_ARRAY_WEIGHTED,
createDoublesBlock(invalidValue),
createDoublesBlock(1.0),
createRleBlock(ImmutableList.of(0.5, 0.75), 1))
.isInstanceOf(TrinoException.class)
.hasErrorCode(INVALID_FUNCTION_ARGUMENT);

// for deprecated approx_percentile with accuracy we only validate accuracy for backward compatibility
assertAggregationFails(
FUNCTION_RESOLUTION,
"approx_percentile",
DOUBLE_APPROXIMATE_PERCENTILE_WEIGHTED_WITH_ACCURACY,
createDoublesBlock(1.0),
createDoublesBlock(1.0),
createDoublesBlock(0.99),
createDoublesBlock(invalidValue))
.hasErrorCode(INVALID_FUNCTION_ARGUMENT);
}
}

private static Block createRleBlock(double percentile, int positionCount)
Expand Down