Skip to content

Commit 17d3c30

Browse files
nielspardonnielspardon
authored
fix(isthmus): std_dev, variance function mappings (#780)
This PR implements the bidirectional conversion between Calcite and Substrait for the statistical aggregate functions standard deviation and variance (`STDDEV_POP`, `STDDEV_SAMP`, `VAR_POP`, `VAR_SAMP`). ### Problem Substrait uses a single function name for both the population and sample variants: - `std_dev` for both `STDDEV_POP` and `STDDEV_SAMP` - `variance` for both `VAR_POP` and `VAR_SAMP` The population (n denominator) vs. sample (n-1 denominator) distinction is captured by a `distribution` value (`POPULATION` or `SAMPLE`). Previously these functions were not mapped, so the existing TPC-DS test cases using them were silently mis-mapped to `AVG` (Calcite represents these statistical functions with `SqlAvgAggFunction`). ### Solution This uses the **non-deprecated** signatures introduced in substrait-io/substrait#1011 (Substrait v0.87.0), which carry the distinction as an **enum function argument** (the leading `distribution` argument, e.g. `std_dev:req_fp64`), rather than the deprecated function-option form (substrait-io/substrait#1019). Resolves #803 Resolves #807 Since the input arguments are cast to `FP64` when necessary, the integer-based signatures proposed in substrait-io/substrait#1012 are not required. **Calcite → Substrait:** - Added function mappings in `FunctionMappings` for all four statistical operators. - `AggregateFunctionConverter` synthesizes the leading `distribution` enum operand based on the Calcite `SqlKind`, so the generic function matcher resolves the enum-arg variant and builds the `EnumArg` automatically (no bespoke option plumbing). - Statistical inputs are cast to `FP64` where necessary. **Substrait → Calcite:** - `FunctionConverter.getSqlOperatorFromSubstraitFunc` disambiguates the population/sample operator from the `distribution` enum argument. - `SubstraitRelNodeConverter` and `PreCalciteAggregateValidator` skip the non-value enum argument when building Calcite aggregate operands. **DSL & shared enum:** - A shared `StatisticalDistribution` enum (in `:core`) is the single source of truth for the `SAMPLE` / `POPULATION` values used by both the DSL builder and isthmus. - `SubstraitBuilder` gains `stddevPopulation`, `stddevSample`, `variancePopulation`, and `varianceSample` convenience methods. ### Non-floating-point inputs `std_dev` / `variance` only define `fp32` / `fp64` signatures, so a statistical aggregate over an integer (or other non-fp) column is rewritten in `SubstraitRelVisitor` to cast the argument to `fp64` (appending a cast column so other aggregates over the same column are unaffected) and cast the result back to the type Calcite inferred. The rewrite is idempotent, so the converted plan is stable under further round trips. ### Testing - `AggregationFunctionsTest` exercises full round trips (POJO ⇄ proto and Substrait ⇄ Calcite) for all four functions, with and without grouping. - `StatisticalFunctionTest` verifies the SQL round trip, asserts that each SQL operator maps to the enum-arg signature (`std_dev:req_fp64` etc.) with the correct `distribution` `EnumArg` and no function options, and covers non-fp (integer) inputs — including a column shared with a non-statistical aggregate, and with grouping. 🤖 Generated with AI --------- Signed-off-by: Niels Pardon <par@zurich.ibm.com>
1 parent f81d841 commit 17d3c30

13 files changed

Lines changed: 896 additions & 34 deletions

‎core/src/main/java/io/substrait/dsl/SubstraitBuilder.java‎

Lines changed: 189 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package io.substrait.dsl;
22

33
import io.substrait.expression.AggregateFunctionInvocation;
4+
import io.substrait.expression.EnumArg;
45
import io.substrait.expression.Expression;
56
import io.substrait.expression.Expression.Cast;
67
import io.substrait.expression.Expression.FailureBehavior;
@@ -13,6 +14,7 @@
1314
import io.substrait.expression.FieldReference;
1415
import io.substrait.expression.FunctionArg;
1516
import io.substrait.expression.FunctionOption;
17+
import io.substrait.expression.StatisticalDistribution;
1618
import io.substrait.expression.WindowBound;
1719
import io.substrait.extension.DefaultExtensionCatalog;
1820
import io.substrait.extension.SimpleExtension;
@@ -1410,6 +1412,193 @@ public Aggregate.Measure sum0(Expression expr) {
14101412
R.I64);
14111413
}
14121414

1415+
/**
1416+
* Creates a population standard deviation aggregate measure for a specific field.
1417+
*
1418+
* <p>Computes the standard deviation using the population formula (n denominator), which
1419+
* considers all values in the dataset as the entire population. This is equivalent to SQL's
1420+
* STDDEV_POP function.
1421+
*
1422+
* @param input the input relation containing the field
1423+
* @param field the zero-based index of the field to aggregate
1424+
* @return an aggregate measure computing population standard deviation with
1425+
* distribution=POPULATION enum argument
1426+
*/
1427+
public Aggregate.Measure stddevPopulation(Rel input, int field) {
1428+
return stddevPopulation(fieldReference(input, field));
1429+
}
1430+
1431+
/**
1432+
* Creates a population standard deviation aggregate measure for an expression.
1433+
*
1434+
* <p>Computes the standard deviation using the population formula (n denominator), which
1435+
* considers all values in the dataset as the entire population. This is equivalent to SQL's
1436+
* STDDEV_POP function.
1437+
*
1438+
* <p>The measure is created with:
1439+
*
1440+
* <ul>
1441+
* <li>Function: Substrait's "std_dev" from the arithmetic extension
1442+
* <li>Argument: distribution=POPULATION (enum argument)
1443+
* <li>Output type: nullable version of the input expression type
1444+
* <li>Aggregation phase: INITIAL_TO_RESULT
1445+
* <li>Invocation: ALL (processes all rows)
1446+
* </ul>
1447+
*
1448+
* @param expr the expression to aggregate (typically a numeric field reference)
1449+
* @return an aggregate measure computing population standard deviation
1450+
*/
1451+
public Aggregate.Measure stddevPopulation(Expression expr) {
1452+
return statisticalAggregate(expr, "std_dev", StatisticalDistribution.POPULATION);
1453+
}
1454+
1455+
/**
1456+
* Creates a sample standard deviation aggregate measure for a specific field.
1457+
*
1458+
* <p>Computes the standard deviation using the sample formula (n-1 denominator), which applies
1459+
* Bessel's correction for sample data. This is equivalent to SQL's STDDEV_SAMP or STDDEV
1460+
* function.
1461+
*
1462+
* @param input the input relation containing the field
1463+
* @param field the zero-based index of the field to aggregate
1464+
* @return an aggregate measure computing sample standard deviation with distribution=SAMPLE enum
1465+
* argument
1466+
*/
1467+
public Aggregate.Measure stddevSample(Rel input, int field) {
1468+
return stddevSample(fieldReference(input, field));
1469+
}
1470+
1471+
/**
1472+
* Creates a sample standard deviation aggregate measure for an expression.
1473+
*
1474+
* <p>Computes the standard deviation using the sample formula (n-1 denominator), which applies
1475+
* Bessel's correction for sample data. This is equivalent to SQL's STDDEV_SAMP or STDDEV
1476+
* function.
1477+
*
1478+
* <p>The measure is created with:
1479+
*
1480+
* <ul>
1481+
* <li>Function: Substrait's "std_dev" from the arithmetic extension
1482+
* <li>Argument: distribution=SAMPLE (enum argument)
1483+
* <li>Output type: nullable version of the input expression type
1484+
* <li>Aggregation phase: INITIAL_TO_RESULT
1485+
* <li>Invocation: ALL (processes all rows)
1486+
* </ul>
1487+
*
1488+
* @param expr the expression to aggregate (typically a numeric field reference)
1489+
* @return an aggregate measure computing sample standard deviation
1490+
*/
1491+
public Aggregate.Measure stddevSample(Expression expr) {
1492+
return statisticalAggregate(expr, "std_dev", StatisticalDistribution.SAMPLE);
1493+
}
1494+
1495+
/**
1496+
* Creates a population variance aggregate measure for a specific field.
1497+
*
1498+
* <p>Computes the variance using the population formula (n denominator), which considers all
1499+
* values in the dataset as the entire population. This is equivalent to SQL's VAR_POP function.
1500+
*
1501+
* @param input the input relation containing the field
1502+
* @param field the zero-based index of the field to aggregate
1503+
* @return an aggregate measure computing population variance with distribution=POPULATION enum
1504+
* argument
1505+
*/
1506+
public Aggregate.Measure variancePopulation(Rel input, int field) {
1507+
return variancePopulation(fieldReference(input, field));
1508+
}
1509+
1510+
/**
1511+
* Creates a population variance aggregate measure for an expression.
1512+
*
1513+
* <p>Computes the variance using the population formula (n denominator), which considers all
1514+
* values in the dataset as the entire population. This is equivalent to SQL's VAR_POP function.
1515+
*
1516+
* <p>The measure is created with:
1517+
*
1518+
* <ul>
1519+
* <li>Function: Substrait's "variance" from the arithmetic extension
1520+
* <li>Argument: distribution=POPULATION (enum argument)
1521+
* <li>Output type: nullable version of the input expression type
1522+
* <li>Aggregation phase: INITIAL_TO_RESULT
1523+
* <li>Invocation: ALL (processes all rows)
1524+
* </ul>
1525+
*
1526+
* @param expr the expression to aggregate (typically a numeric field reference)
1527+
* @return an aggregate measure computing population variance
1528+
*/
1529+
public Aggregate.Measure variancePopulation(Expression expr) {
1530+
return statisticalAggregate(expr, "variance", StatisticalDistribution.POPULATION);
1531+
}
1532+
1533+
/**
1534+
* Creates a sample variance aggregate measure for a specific field.
1535+
*
1536+
* <p>Computes the variance using the sample formula (n-1 denominator), which applies Bessel's
1537+
* correction for sample data. This is equivalent to SQL's VAR_SAMP or VARIANCE function.
1538+
*
1539+
* @param input the input relation containing the field
1540+
* @param field the zero-based index of the field to aggregate
1541+
* @return an aggregate measure computing sample variance with distribution=SAMPLE enum argument
1542+
*/
1543+
public Aggregate.Measure varianceSample(Rel input, int field) {
1544+
return varianceSample(fieldReference(input, field));
1545+
}
1546+
1547+
/**
1548+
* Creates a sample variance aggregate measure for an expression.
1549+
*
1550+
* <p>Computes the variance using the sample formula (n-1 denominator), which applies Bessel's
1551+
* correction for sample data. This is equivalent to SQL's VAR_SAMP or VARIANCE function.
1552+
*
1553+
* <p>The measure is created with:
1554+
*
1555+
* <ul>
1556+
* <li>Function: Substrait's "variance" from the arithmetic extension
1557+
* <li>Argument: distribution=SAMPLE (enum argument)
1558+
* <li>Output type: nullable version of the input expression type
1559+
* <li>Aggregation phase: INITIAL_TO_RESULT
1560+
* <li>Invocation: ALL (processes all rows)
1561+
* </ul>
1562+
*
1563+
* @param expr the expression to aggregate (typically a numeric field reference)
1564+
* @return an aggregate measure computing sample variance
1565+
*/
1566+
public Aggregate.Measure varianceSample(Expression expr) {
1567+
return statisticalAggregate(expr, "variance", StatisticalDistribution.SAMPLE);
1568+
}
1569+
1570+
/**
1571+
* Helper method to create statistical aggregate measures (std_dev, variance) with a {@code
1572+
* distribution} enum argument.
1573+
*
1574+
* <p>Uses the non-deprecated function signatures that carry the population/sample distinction as
1575+
* a leading {@code distribution} {@link EnumArg} (e.g. {@code std_dev:req_fp64}).
1576+
*
1577+
* @param expr the expression to aggregate
1578+
* @param functionName the Substrait function name ("std_dev" or "variance")
1579+
* @param distribution the distribution type (SAMPLE or POPULATION)
1580+
* @return an aggregate measure with the specified distribution argument
1581+
*/
1582+
private Aggregate.Measure statisticalAggregate(
1583+
Expression expr, String functionName, StatisticalDistribution distribution) {
1584+
String typeString = ToTypeString.apply(expr.getType());
1585+
SimpleExtension.AggregateFunctionVariant declaration =
1586+
extensions.getAggregateFunction(
1587+
SimpleExtension.FunctionAnchor.of(
1588+
DefaultExtensionCatalog.FUNCTIONS_ARITHMETIC,
1589+
String.format("%s:req_%s", functionName, typeString)));
1590+
EnumArg distributionArg =
1591+
EnumArg.of((SimpleExtension.EnumArgument) declaration.args().get(0), distribution.name());
1592+
return measure(
1593+
AggregateFunctionInvocation.builder()
1594+
.arguments(Arrays.asList(distributionArg, expr))
1595+
.outputType(TypeCreator.asNullable(expr.getType()))
1596+
.declaration(declaration)
1597+
.aggregationPhase(Expression.AggregationPhase.INITIAL_TO_RESULT)
1598+
.invocation(Expression.AggregationInvocation.ALL)
1599+
.build());
1600+
}
1601+
14131602
private Aggregate.Measure singleArgumentArithmeticAggregate(
14141603
Expression expr, String functionName, Type outputType) {
14151604
String typeString = ToTypeString.apply(expr.getType());

‎core/src/main/java/io/substrait/expression/StatisticalDistribution.java‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
package io.substrait.expression;
2+
3+
/**
4+
* The {@code distribution} enum argument of the Substrait {@code std_dev} and {@code variance}
5+
* aggregate functions.
6+
*
7+
* <p>Distinguishes between the sample (n-1 denominator, Bessel's correction) and population (n
8+
* denominator) variants. The enum constant names match the Substrait extension's enum option names
9+
* ({@code SAMPLE} / {@code POPULATION}), so {@link #name()} yields the value used to build an
10+
* {@link EnumArg}.
11+
*/
12+
public enum StatisticalDistribution {
13+
/** Sample distribution (uses the n-1 denominator, Bessel's correction). */
14+
SAMPLE,
15+
/** Population distribution (uses the n denominator). */
16+
POPULATION
17+
}

‎isthmus/src/main/java/io/substrait/isthmus/AggregateFunctions.java‎

Lines changed: 52 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,30 @@ public class AggregateFunctions {
2929
/** Substrait-specific AVG aggregate function (nullable return type). */
3030
public static final SqlAggFunction AVG = new SubstraitAvgAggFunction(SqlKind.AVG);
3131

32+
/**
33+
* Standard deviation (population) aggregate function. Maps to Substrait's std_dev function with
34+
* distribution=POPULATION enum argument.
35+
*/
36+
public static SqlAggFunction STDDEV_POP = new SubstraitAvgAggFunction(SqlKind.STDDEV_POP);
37+
38+
/**
39+
* Standard deviation (sample) aggregate function. Maps to Substrait's std_dev function with
40+
* distribution=SAMPLE enum argument.
41+
*/
42+
public static SqlAggFunction STDDEV_SAMP = new SubstraitAvgAggFunction(SqlKind.STDDEV_SAMP);
43+
44+
/**
45+
* Variance (population) aggregate function. Maps to Substrait's variance function with
46+
* distribution=POPULATION enum argument.
47+
*/
48+
public static SqlAggFunction VAR_POP = new SubstraitAvgAggFunction(SqlKind.VAR_POP);
49+
50+
/**
51+
* Variance (sample) aggregate function. Maps to Substrait's variance function with
52+
* distribution=SAMPLE enum argument.
53+
*/
54+
public static SqlAggFunction VAR_SAMP = new SubstraitAvgAggFunction(SqlKind.VAR_SAMP);
55+
3256
/** Substrait-specific SUM aggregate function (nullable return type). */
3357
public static final SqlAggFunction SUM = new SubstraitSumAggFunction();
3458

@@ -42,18 +66,34 @@ public class AggregateFunctions {
4266
* @return optional containing Substrait equivalent if conversion applies
4367
*/
4468
public static Optional<SqlAggFunction> toSubstraitAggVariant(SqlAggFunction aggFunction) {
45-
if (aggFunction instanceof SqlMinMaxAggFunction) {
46-
SqlMinMaxAggFunction fun = (SqlMinMaxAggFunction) aggFunction;
47-
return Optional.of(
48-
fun.getKind() == SqlKind.MIN ? AggregateFunctions.MIN : AggregateFunctions.MAX);
49-
} else if (aggFunction instanceof SqlAvgAggFunction) {
50-
return Optional.of(AggregateFunctions.AVG);
51-
} else if (aggFunction instanceof SqlSumAggFunction) {
52-
return Optional.of(AggregateFunctions.SUM);
53-
} else if (aggFunction instanceof SqlSumEmptyIsZeroAggFunction) {
54-
return Optional.of(AggregateFunctions.SUM0);
55-
} else {
56-
return Optional.empty();
69+
// First check by SqlKind to handle all statistical functions
70+
SqlKind kind = aggFunction.getKind();
71+
switch (kind) {
72+
case MIN:
73+
return Optional.of(AggregateFunctions.MIN);
74+
case MAX:
75+
return Optional.of(AggregateFunctions.MAX);
76+
case AVG:
77+
return Optional.of(AggregateFunctions.AVG);
78+
case STDDEV_POP:
79+
return Optional.of(AggregateFunctions.STDDEV_POP);
80+
case STDDEV_SAMP:
81+
return Optional.of(AggregateFunctions.STDDEV_SAMP);
82+
case VAR_POP:
83+
return Optional.of(AggregateFunctions.VAR_POP);
84+
case VAR_SAMP:
85+
return Optional.of(AggregateFunctions.VAR_SAMP);
86+
case SUM:
87+
case SUM0:
88+
// Check instance type for SUM variants
89+
if (aggFunction instanceof SqlSumEmptyIsZeroAggFunction) {
90+
return Optional.of(AggregateFunctions.SUM0);
91+
} else if (aggFunction instanceof SqlSumAggFunction) {
92+
return Optional.of(AggregateFunctions.SUM);
93+
}
94+
return Optional.empty();
95+
default:
96+
return Optional.empty();
5797
}
5898
}
5999

‎isthmus/src/main/java/io/substrait/isthmus/PreCalciteAggregateValidator.java‎

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -46,8 +46,11 @@ public static boolean isValidCalciteAggregate(Aggregate aggregate) {
4646
*/
4747
private static boolean isValidCalciteMeasure(Aggregate.Measure measure) {
4848
return
49-
// all function arguments to measures must be field references
50-
measure.getFunction().arguments().stream().allMatch(farg -> isSimpleFieldReference(farg))
49+
// all value (Expression) function arguments to measures must be field references; non-value
50+
// arguments such as the std_dev/variance "distribution" enum argument are exempt
51+
measure.getFunction().arguments().stream()
52+
.filter(farg -> farg instanceof Expression)
53+
.allMatch(farg -> isSimpleFieldReference(farg))
5154
&&
5255
// all sort fields must be field references
5356
measure.getFunction().sort().stream().allMatch(sf -> isSimpleFieldReference(sf.expr()))
@@ -157,9 +160,9 @@ public static Aggregate transformToValidCalciteAggregate(Aggregate aggregate) {
157160
private Aggregate.Measure updateMeasure(Aggregate.Measure measure) {
158161
AggregateFunctionInvocation oldAggregateFunctionInvocation = measure.getFunction();
159162

160-
List<Expression> newFunctionArgs =
163+
List<FunctionArg> newFunctionArgs =
161164
oldAggregateFunctionInvocation.arguments().stream()
162-
.map(this::projectOutNonFieldReference)
165+
.map(this::projectOutNonFieldReferenceArg)
163166
.collect(Collectors.toList());
164167

165168
List<Expression.SortField> newSortFields =
@@ -194,11 +197,13 @@ private Aggregate.Grouping updateGrouping(Aggregate.Grouping grouping) {
194197
return Aggregate.Grouping.builder().expressions(newGroupingExpressions).build();
195198
}
196199

197-
private Expression projectOutNonFieldReference(FunctionArg farg) {
200+
private FunctionArg projectOutNonFieldReferenceArg(FunctionArg farg) {
198201
if ((farg instanceof Expression)) {
199202
return projectOutNonFieldReference((Expression) farg);
200203
} else {
201-
throw new IllegalArgumentException("cannot handle non-expression argument for aggregate");
204+
// Non-value arguments (e.g. the std_dev/variance "distribution" enum argument) are not
205+
// field references and are passed through unchanged.
206+
return farg;
202207
}
203208
}
204209

‎isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -385,8 +385,11 @@ public RelNode visit(Aggregate aggregate, Context context) throws RuntimeExcepti
385385

386386
private AggregateCall fromMeasure(Aggregate.Measure measure, Context context) {
387387
List<FunctionArg> eArgs = measure.getFunction().arguments();
388+
// Only value (Expression) arguments map to Calcite aggregate operands. Enum arguments such as
389+
// the std_dev/variance "distribution" are used to disambiguate the operator, not as operands.
388390
List<RexNode> arguments =
389-
IntStream.range(0, measure.getFunction().arguments().size())
391+
IntStream.range(0, eArgs.size())
392+
.filter(i -> eArgs.get(i) instanceof Expression)
390393
.mapToObj(
391394
i ->
392395
eArgs
@@ -399,7 +402,9 @@ private AggregateCall fromMeasure(Aggregate.Measure measure, Context context) {
399402
.collect(java.util.stream.Collectors.toList());
400403
Optional<SqlOperator> operator =
401404
aggregateFunctionConverter.getSqlOperatorFromSubstraitFunc(
402-
measure.getFunction().declaration().key(), measure.getFunction().outputType());
405+
measure.getFunction().declaration().key(),
406+
measure.getFunction().outputType(),
407+
measure.getFunction().arguments());
403408
if (!operator.isPresent()) {
404409
throw new IllegalArgumentException(
405410
String.format(

0 commit comments

Comments
 (0)