Skip to content

Commit de95ffb

Browse files
ahkcsahkcs
authored
Detect integral SUM overflow and fix BIGINT AVG (#5612)
1 parent e905aec commit de95ffb

63 files changed

Lines changed: 966 additions & 157 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎api/src/test/java/org/opensearch/sql/api/UnifiedQueryPlannerSqlV2Test.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -286,7 +286,7 @@ SELECT department, SUM(age) AS total FROM catalog.employees GROUP BY department
286286
.assertPlan(
287287
"""
288288
LogicalProject(department=[$0], total=[$1])
289-
LogicalAggregate(group=[{0}], SUM(age)=[SUM($1)])
289+
LogicalAggregate(group=[{0}], SUM(age)=[CHECKED_LONG_SUM($1)])
290290
LogicalProject(department=[$3], age=[$2])
291291
LogicalTableScan(table=[[catalog, employees]])
292292
""");
@@ -366,7 +366,7 @@ SELECT department, SUM(age) FILTER(WHERE age > 30) FROM catalog.employees
366366
""")
367367
.assertPlan(
368368
"""
369-
LogicalAggregate(group=[{0}], SUM(age) FILTER(WHERE age > 30)=[SUM($1) FILTER $2])
369+
LogicalAggregate(group=[{0}], SUM(age) FILTER(WHERE age > 30)=[CHECKED_LONG_SUM($1) FILTER $2])
370370
LogicalProject(department=[$3], age=[$2], $f3=[>($2, 30)])
371371
LogicalTableScan(table=[[catalog, employees]])
372372
""");
@@ -487,7 +487,7 @@ SELECT name, SUM(age) OVER(PARTITION BY department ORDER BY age) FROM catalog.em
487487
""")
488488
.assertPlan(
489489
"""
490-
LogicalProject(name=[$1], SUM(age) OVER(PARTITION BY department ORDER BY age)=[SUM($2) OVER (PARTITION BY $3 ORDER BY $2 NULLS FIRST)])
490+
LogicalProject(name=[$1], SUM(age) OVER(PARTITION BY department ORDER BY age)=[CHECKED_LONG_SUM($2) OVER (PARTITION BY $3 ORDER BY $2 NULLS FIRST)])
491491
LogicalTableScan(table=[[catalog, employees]])
492492
""");
493493
}

‎core/src/main/java/org/opensearch/sql/calcite/udf/udaf/BigintAvgAggFunction.java‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.calcite.udf.udaf;
7+
8+
/** BIGINT average aggregate that accumulates in double to avoid an intermediate long overflow. */
9+
public class BigintAvgAggFunction {
10+
11+
public static Accumulator init() {
12+
return new Accumulator();
13+
}
14+
15+
public static Accumulator add(Accumulator accumulator, Long value) {
16+
if (value != null) {
17+
accumulator.sum += value;
18+
accumulator.count++;
19+
}
20+
return accumulator;
21+
}
22+
23+
public static Double result(Accumulator accumulator) {
24+
return accumulator.count == 0 ? null : accumulator.sum / accumulator.count;
25+
}
26+
27+
public static class Accumulator {
28+
private double sum;
29+
private long count;
30+
}
31+
}

‎core/src/main/java/org/opensearch/sql/calcite/udf/udaf/CheckedLongSumAggFunction.java‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.calcite.udf.udaf;
7+
8+
/** BIGINT sum aggregate that throws when its running long accumulator overflows. */
9+
public class CheckedLongSumAggFunction {
10+
11+
public static long init() {
12+
return 0L;
13+
}
14+
15+
public static long add(long accumulator, long value) {
16+
return Math.addExact(accumulator, value);
17+
}
18+
19+
public static long result(long accumulator) {
20+
return accumulator;
21+
}
22+
}

‎core/src/main/java/org/opensearch/sql/calcite/utils/UserDefinedFunctionUtils.java‎

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -94,9 +94,32 @@ public static SqlUserDefinedAggFunction createUserDefinedAggFunction(
9494
String functionName,
9595
SqlReturnTypeInference returnType,
9696
@Nullable UDFOperandMetadata operandMetadata) {
97+
return createReflectiveAggFunction(udafClass, functionName, returnType, operandMetadata);
98+
}
99+
100+
/** Creates an aggregate function from a class following Calcite's reflective UDAF convention. */
101+
public static SqlUserDefinedAggFunction createReflectiveAggFunction(
102+
Class<?> udafClass,
103+
String functionName,
104+
SqlReturnTypeInference returnType,
105+
@Nullable UDFOperandMetadata operandMetadata) {
106+
return createReflectiveAggFunction(
107+
udafClass, functionName, SqlKind.OTHER_FUNCTION, returnType, operandMetadata);
108+
}
109+
110+
/**
111+
* Creates an aggregate function whose kind remains visible to planner rules while execution uses
112+
* the supplied reflective UDAF.
113+
*/
114+
public static SqlUserDefinedAggFunction createReflectiveAggFunction(
115+
Class<?> udafClass,
116+
String functionName,
117+
SqlKind kind,
118+
SqlReturnTypeInference returnType,
119+
@Nullable UDFOperandMetadata operandMetadata) {
97120
return new SqlUserDefinedAggFunction(
98121
new SqlIdentifier(functionName, SqlParserPos.ZERO),
99-
SqlKind.OTHER_FUNCTION,
122+
kind,
100123
returnType,
101124
null,
102125
operandMetadata,

‎core/src/main/java/org/opensearch/sql/expression/function/PPLBuiltinOperators.java‎

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import static org.opensearch.sql.calcite.utils.UserDefinedFunctionUtils.adaptExprMethodToUDF;
99
import static org.opensearch.sql.calcite.utils.UserDefinedFunctionUtils.adaptExprMethodWithPropertiesToUDF;
1010
import static org.opensearch.sql.calcite.utils.UserDefinedFunctionUtils.adaptMathFunctionToUDF;
11+
import static org.opensearch.sql.calcite.utils.UserDefinedFunctionUtils.createReflectiveAggFunction;
1112
import static org.opensearch.sql.calcite.utils.UserDefinedFunctionUtils.createUserDefinedAggFunction;
1213

1314
import com.google.common.base.Suppliers;
@@ -29,6 +30,8 @@
2930
import org.apache.calcite.sql.type.SqlTypeTransforms;
3031
import org.apache.calcite.sql.util.ReflectiveSqlOperatorTable;
3132
import org.apache.calcite.util.BuiltInMethod;
33+
import org.opensearch.sql.calcite.udf.udaf.BigintAvgAggFunction;
34+
import org.opensearch.sql.calcite.udf.udaf.CheckedLongSumAggFunction;
3235
import org.opensearch.sql.calcite.udf.udaf.DistinctCountApproxLogicalAggFunction;
3336
import org.opensearch.sql.calcite.udf.udaf.FirstAggFunction;
3437
import org.opensearch.sql.calcite.udf.udaf.LastAggFunction;
@@ -449,7 +452,6 @@ public class PPLBuiltinOperators extends ReflectiveSqlOperatorTable {
449452
new NumberToStringFunction().toUDF("NUMBER_TO_STRING");
450453
public static final SqlOperator TONUMBER = new ToNumberFunction().toUDF("TONUMBER");
451454
public static final SqlOperator TOSTRING = new ToStringFunction().toUDF("TOSTRING");
452-
453455
// PPL Convert command functions
454456
public static final SqlOperator AUTO = new AutoConvertFunction().toUDF("AUTO");
455457
public static final SqlOperator NUM = new NumConvertFunction().toUDF("NUM");
@@ -488,6 +490,20 @@ public class PPLBuiltinOperators extends ReflectiveSqlOperatorTable {
488490
new NullableSqlAvgAggFunction(SqlKind.VAR_POP);
489491
public static final SqlAggFunction VAR_SAMP_NULLABLE =
490492
new NullableSqlAvgAggFunction(SqlKind.VAR_SAMP);
493+
public static final SqlAggFunction CHECKED_LONG_SUM =
494+
createReflectiveAggFunction(
495+
CheckedLongSumAggFunction.class,
496+
"CHECKED_LONG_SUM",
497+
SqlKind.SUM,
498+
ReturnTypes.BIGINT_FORCE_NULLABLE,
499+
PPLOperandTypes.NUMERIC);
500+
public static final SqlAggFunction BIGINT_AVG =
501+
createReflectiveAggFunction(
502+
BigintAvgAggFunction.class,
503+
"AVG",
504+
SqlKind.AVG,
505+
ReturnTypes.DOUBLE_NULLABLE,
506+
PPLOperandTypes.NUMERIC);
491507
public static final SqlAggFunction TAKE =
492508
createUserDefinedAggFunction(
493509
TakeAggFunction.class,

‎core/src/main/java/org/opensearch/sql/expression/function/PPLFuncImpTable.java‎

Lines changed: 34 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -283,6 +283,7 @@
283283
import java.util.Optional;
284284
import java.util.StringJoiner;
285285
import java.util.concurrent.ConcurrentHashMap;
286+
import java.util.function.Function;
286287
import java.util.regex.Pattern;
287288
import java.util.regex.PatternSyntaxException;
288289
import java.util.stream.Collectors;
@@ -1595,23 +1596,45 @@ void register(
15951596
}
15961597

15971598
void registerOperator(BuiltinFunctionName functionName, SqlAggFunction aggFunction) {
1598-
SqlOperandTypeChecker innerTypeChecker = extractTypeCheckerFromUDF(aggFunction);
1599+
registerOperator(functionName, aggFunction, field -> aggFunction);
1600+
}
1601+
1602+
void registerOperator(
1603+
BuiltinFunctionName functionName,
1604+
SqlAggFunction typeCheckerSource,
1605+
Function<RexNode, SqlAggFunction> aggFunctionSelector) {
1606+
SqlOperandTypeChecker innerTypeChecker = extractTypeCheckerFromUDF(typeCheckerSource);
15991607
PPLTypeChecker typeChecker =
16001608
wrapSqlOperandTypeChecker(innerTypeChecker, functionName.name(), true);
16011609
AggHandler handler =
16021610
(distinct, field, argList, ctx) -> {
16031611
List<RexNode> newArgList =
16041612
argList.stream().map(PlanUtils::derefMapCall).collect(Collectors.toList());
16051613
return UserDefinedFunctionUtils.makeAggregateCall(
1606-
aggFunction, List.of(field), newArgList, ctx.relBuilder);
1614+
aggFunctionSelector.apply(field), List.of(field), newArgList, ctx.relBuilder);
16071615
};
16081616
register(functionName, handler, typeChecker);
16091617
}
16101618

1619+
/** Registers checked integral sums while retaining standard SUM behavior for other types. */
1620+
void registerSumOperator() {
1621+
registerOperator(
1622+
SUM,
1623+
SqlStdOperatorTable.SUM,
1624+
field ->
1625+
isIntegral(field.getType().getSqlTypeName())
1626+
? PPLBuiltinOperators.CHECKED_LONG_SUM
1627+
: SqlStdOperatorTable.SUM);
1628+
}
1629+
1630+
private static boolean isIntegral(SqlTypeName typeName) {
1631+
return SqlTypeName.INT_TYPES.contains(typeName);
1632+
}
1633+
16111634
void populate() {
16121635
registerOperator(MAX, SqlStdOperatorTable.MAX);
16131636
registerOperator(MIN, SqlStdOperatorTable.MIN);
1614-
registerOperator(SUM, SqlStdOperatorTable.SUM);
1637+
registerSumOperator();
16151638
registerOperator(VARSAMP, PPLBuiltinOperators.VAR_SAMP_NULLABLE);
16161639
registerOperator(VARPOP, PPLBuiltinOperators.VAR_POP_NULLABLE);
16171640
registerOperator(STDDEV_SAMP, PPLBuiltinOperators.STDDEV_SAMP_NULLABLE);
@@ -1630,7 +1653,14 @@ void populate() {
16301653

16311654
register(
16321655
AVG,
1633-
(distinct, field, argList, ctx) -> ctx.relBuilder.avg(distinct, null, field),
1656+
(distinct, field, argList, ctx) -> {
1657+
if (field.getType().getSqlTypeName() == SqlTypeName.BIGINT) {
1658+
return ctx.relBuilder
1659+
.aggregateCall(PPLBuiltinOperators.BIGINT_AVG, field)
1660+
.distinct(distinct);
1661+
}
1662+
return ctx.relBuilder.avg(distinct, null, field);
1663+
},
16341664
wrapSqlOperandTypeChecker(
16351665
SqlStdOperatorTable.AVG.getOperandTypeChecker(), AVG.name(), false));
16361666

‎core/src/test/java/org/opensearch/sql/calcite/udf/udaf/BigintAvgAggFunctionTest.java‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.calcite.udf.udaf;
7+
8+
import static org.junit.jupiter.api.Assertions.assertEquals;
9+
import static org.junit.jupiter.api.Assertions.assertNull;
10+
11+
import org.apache.calcite.sql.SqlKind;
12+
import org.junit.jupiter.api.Test;
13+
import org.opensearch.sql.expression.function.PPLBuiltinOperators;
14+
15+
class BigintAvgAggFunctionTest {
16+
17+
@Test
18+
void retainsAvgKindForPushdownRules() {
19+
assertEquals(SqlKind.AVG, PPLBuiltinOperators.BIGINT_AVG.getKind());
20+
}
21+
22+
@Test
23+
void averagesWithoutLongOverflow() {
24+
BigintAvgAggFunction.Accumulator accumulator = BigintAvgAggFunction.init();
25+
accumulator = BigintAvgAggFunction.add(accumulator, Long.MAX_VALUE);
26+
accumulator = BigintAvgAggFunction.add(accumulator, Long.MAX_VALUE);
27+
28+
assertEquals((double) Long.MAX_VALUE, BigintAvgAggFunction.result(accumulator));
29+
}
30+
31+
@Test
32+
void ignoresNulls() {
33+
BigintAvgAggFunction.Accumulator accumulator = BigintAvgAggFunction.init();
34+
accumulator = BigintAvgAggFunction.add(accumulator, null);
35+
accumulator = BigintAvgAggFunction.add(accumulator, 10L);
36+
37+
assertEquals(10D, BigintAvgAggFunction.result(accumulator));
38+
}
39+
40+
@Test
41+
void returnsNullForEmptyInput() {
42+
assertNull(BigintAvgAggFunction.result(BigintAvgAggFunction.init()));
43+
}
44+
}

‎core/src/test/java/org/opensearch/sql/calcite/udf/udaf/CheckedLongSumAggFunctionTest.java‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.calcite.udf.udaf;
7+
8+
import static org.junit.jupiter.api.Assertions.assertEquals;
9+
import static org.junit.jupiter.api.Assertions.assertThrows;
10+
11+
import org.apache.calcite.sql.SqlKind;
12+
import org.junit.jupiter.api.Test;
13+
import org.opensearch.sql.expression.function.PPLBuiltinOperators;
14+
15+
class CheckedLongSumAggFunctionTest {
16+
17+
@Test
18+
void retainsSumKindForPlannerRules() {
19+
assertEquals(SqlKind.SUM, PPLBuiltinOperators.CHECKED_LONG_SUM.getKind());
20+
}
21+
22+
@Test
23+
void sumsExactly() {
24+
long accumulator = CheckedLongSumAggFunction.init();
25+
accumulator = CheckedLongSumAggFunction.add(accumulator, 1L << 62);
26+
accumulator = CheckedLongSumAggFunction.add(accumulator, 1L);
27+
28+
assertEquals((1L << 62) + 1L, CheckedLongSumAggFunction.result(accumulator));
29+
}
30+
31+
@Test
32+
void throwsOnPositiveOverflow() {
33+
assertThrows(
34+
ArithmeticException.class, () -> CheckedLongSumAggFunction.add(Long.MAX_VALUE, 1L));
35+
}
36+
37+
@Test
38+
void throwsOnNegativeOverflow() {
39+
assertThrows(
40+
ArithmeticException.class, () -> CheckedLongSumAggFunction.add(Long.MIN_VALUE, -1L));
41+
}
42+
}

‎docs/user/ppl/functions/expressions.md‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,14 @@ Arithmetic expressions are formed by combining numeric literals and binary arith
1313

1414
### Overflow behavior
1515

16-
Integer and long arithmetic operations (`+`, `-`, `*`) in `eval` expressions detect overflow and return an error instead of silently wrapping. For example, `eval x = int_field + 1` where `int_field` is `2147483647` (integer max) returns an error rather than `-2147483648`. Floating-point (`float`, `double`) arithmetic follows IEEE 754 and does not produce overflow errors.
16+
Long (`BIGINT`) arithmetic operations (`+`, `-`, `*`) in `eval` expressions detect overflow and return an error instead of silently wrapping. Narrower integer operands are widened before arithmetic, so crossing the 32-bit integer boundary does not overflow. Floating-point (`float`, `double`) arithmetic follows IEEE 754 and does not produce overflow errors.
17+
18+
The accumulator used by `stats sum(integral_field)` depends on whether Calcite pushdown is enabled:
19+
20+
- With `plugins.calcite.pushdown.enabled=true` (the default), OpenSearch uses its native double-based `sum`, then the result is checked and narrowed to `BIGINT`. Large in-range sums can lose low-order precision. A small overflow that rounds to a signed `BIGINT` boundary can also be indistinguishable from an in-range sum and may saturate at the boundary instead of returning an error.
21+
- With `plugins.calcite.pushdown.enabled=false`, Calcite uses an exact `BIGINT` accumulator and `Math.addExact` for every addition. It returns an error as soon as the running sum exceeds the `BIGINT` range.
22+
23+
For example, summing `4611686018427387904` (`2^62`) and `1` returns the exact `4611686018427387905` without pushdown. With pushdown enabled, the double accumulator cannot represent the low-order `1`, so the result is `4611686018427387904`.
1724

1825
### Precedence
1926

@@ -189,4 +196,3 @@ fetched rows / total rows = 2/2
189196
| 28 |
190197
+-----+
191198
```
192-

0 commit comments

Comments
 (0)