diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/transform/BeamBuiltinAggregations.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/transform/BeamBuiltinAggregations.java index 3fc299bd5a33..2800edfbb99a 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/transform/BeamBuiltinAggregations.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/transform/BeamBuiltinAggregations.java @@ -83,6 +83,8 @@ public class BeamBuiltinAggregations { typeName -> new DropNullFn(BeamBuiltinAggregations.createBitAnd(typeName))) .put("VAR_POP", t -> VarianceFn.newPopulation(t.getTypeName())) .put("VAR_SAMP", t -> VarianceFn.newSample(t.getTypeName())) + .put("STDDEV_POP", t -> VarianceFn.newPopulationStddev(t.getTypeName())) + .put("STDDEV_SAMP", t -> VarianceFn.newSampleStddev(t.getTypeName())) .put("COVAR_POP", t -> CovarianceFn.newPopulation(t.getTypeName())) .put("COVAR_SAMP", t -> CovarianceFn.newSample(t.getTypeName())) .put("COUNTIF", typeName -> CountIf.combineFn()) diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/transform/agg/VarianceFn.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/transform/agg/VarianceFn.java index dd2cd3b20952..906bac7add52 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/transform/agg/VarianceFn.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/transform/agg/VarianceFn.java @@ -75,8 +75,12 @@ public class VarianceFn extends Combine.CombineFn decimalConverter; + private final boolean isSample; // flag to determine return value should be Variance Pop or Sample + // When true, extractOutput returns the square root of the variance (i.e. standard deviation). + // Beam's enumerable bridge cannot translate a SQRT call layered on top of a window VAR_SAMP, so + // STDDEV_SAMP / STDDEV_POP are computed end-to-end inside this combiner instead. + private final boolean isStddev; + private final SerializableFunction decimalConverter; public static VarianceFn newPopulation(Schema.TypeName typeName) { return newPopulation(BigDecimalConverter.forSqlType(typeName)); @@ -85,7 +89,7 @@ public static VarianceFn newPopulation(Schema.TypeName typeName) { public static VarianceFn newPopulation( SerializableFunction decimalConverter) { - return new VarianceFn<>(POP, decimalConverter); + return new VarianceFn<>(POP, false, decimalConverter); } public static VarianceFn newSample(Schema.TypeName typeName) { @@ -95,11 +99,21 @@ public static VarianceFn newSample(Schema.TypeName typeName) { public static VarianceFn newSample( SerializableFunction decimalConverter) { - return new VarianceFn<>(SAMPLE, decimalConverter); + return new VarianceFn<>(SAMPLE, false, decimalConverter); } - private VarianceFn(boolean isSample, SerializableFunction decimalConverter) { + public static VarianceFn newSampleStddev(Schema.TypeName typeName) { + return new VarianceFn<>(SAMPLE, true, BigDecimalConverter.forSqlType(typeName)); + } + + public static VarianceFn newPopulationStddev(Schema.TypeName typeName) { + return new VarianceFn<>(POP, true, BigDecimalConverter.forSqlType(typeName)); + } + + private VarianceFn( + boolean isSample, boolean isStddev, SerializableFunction decimalConverter) { this.isSample = isSample; + this.isStddev = isStddev; this.decimalConverter = decimalConverter; } @@ -133,7 +147,19 @@ public Coder getAccumulatorCoder( @Override public T extractOutput(VarianceAccumulator accumulator) { - return decimalConverter.apply(getVariance(accumulator)); + BigDecimal result = getVariance(accumulator); + if (result != null && isStddev) { + double doubleVal = result.doubleValue(); + if (doubleVal < 0.0) { + doubleVal = 0.0; // Clamp negative variance due to numerical instability + } + double sqrtVal = Math.sqrt(doubleVal); + if (Double.isInfinite(sqrtVal)) { + return decimalConverter.apply(result.sqrt(MATH_CTX)); + } + result = BigDecimal.valueOf(sqrtVal); + } + return decimalConverter.apply(result); } private BigDecimal getVariance(VarianceAccumulator variance) { diff --git a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlDslAggregationVarianceTest.java b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlDslAggregationVarianceTest.java index 808b27aaac4c..e2c548acf718 100644 --- a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlDslAggregationVarianceTest.java +++ b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlDslAggregationVarianceTest.java @@ -30,7 +30,10 @@ import org.junit.Rule; import org.junit.Test; -/** Integration tests for {@code VAR_POP} and {@code VAR_SAMP}. */ +/** + * Integration tests for {@code VAR_POP}, {@code VAR_SAMP}, {@code STDDEV_POP} and {@code + * STDDEV_SAMP}. + */ public class BeamSqlDslAggregationVarianceTest { private static final double PRECISION = 1e-7; @@ -94,4 +97,42 @@ public void testSampleVarianceInt() { pipeline.run().waitUntilFinish(); } + + @Test + public void testPopulationStddevDouble() { + String sql = "SELECT STDDEV_POP(f_double) FROM PCOLLECTION GROUP BY f_int2"; + + PAssert.that(boundedInput.apply(SqlTransform.query(sql))) + .satisfies(matchesScalar(5.138887357, PRECISION)); + + pipeline.run().waitUntilFinish(); + } + + @Test + public void testPopulationStddevInt() { + String sql = "SELECT STDDEV_POP(f_int) FROM PCOLLECTION GROUP BY f_int2"; + + PAssert.that(boundedInput.apply(SqlTransform.query(sql))).satisfies(matchesScalar(5)); + + pipeline.run().waitUntilFinish(); + } + + @Test + public void testSampleStddevDouble() { + String sql = "SELECT STDDEV_SAMP(f_double) FROM PCOLLECTION GROUP BY f_int2"; + + PAssert.that(boundedInput.apply(SqlTransform.query(sql))) + .satisfies(matchesScalar(5.550632739, PRECISION)); + + pipeline.run().waitUntilFinish(); + } + + @Test + public void testSampleStddevInt() { + String sql = "SELECT STDDEV_SAMP(f_int) FROM PCOLLECTION GROUP BY f_int2"; + + PAssert.that(boundedInput.apply(SqlTransform.query(sql))).satisfies(matchesScalar(5)); + + pipeline.run().waitUntilFinish(); + } } diff --git a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/transform/agg/VarianceFnTest.java b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/transform/agg/VarianceFnTest.java index f7a8ad1fa06b..0671a3caaa68 100644 --- a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/transform/agg/VarianceFnTest.java +++ b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/transform/agg/VarianceFnTest.java @@ -26,6 +26,7 @@ import java.util.Arrays; import org.apache.beam.sdk.coders.CoderRegistry; import org.apache.beam.sdk.coders.VarIntCoder; +import org.apache.beam.sdk.schemas.Schema; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.Parameterized; @@ -51,18 +52,38 @@ public static Iterable varianceFns() { VarianceFn.newSample(BigDecimal::intValue), newVarianceAccumulator(FIFTEEN, FOUR, ZERO), 5 + }, + { + VarianceFn.newPopulationStddev(Schema.TypeName.INT32), + newVarianceAccumulator(new BigDecimal(36), new BigDecimal(4), ZERO), + 3 + }, + { + VarianceFn.newSampleStddev(Schema.TypeName.INT32), + newVarianceAccumulator(new BigDecimal(36), new BigDecimal(5), ZERO), + 3 + }, + { + VarianceFn.newPopulationStddev(Schema.TypeName.DOUBLE), + newVarianceAccumulator(new BigDecimal("1e700"), BigDecimal.ONE, ZERO), + Double.POSITIVE_INFINITY + }, + { + VarianceFn.newPopulationStddev(Schema.TypeName.FLOAT), + newVarianceAccumulator(new BigDecimal("1e700"), BigDecimal.ONE, ZERO), + Float.POSITIVE_INFINITY } }); } private VarianceFn varianceFn; private VarianceAccumulator testAccumulatorInput; - private int expectedExtractedResult; + private Object expectedExtractedResult; public VarianceFnTest( VarianceFn varianceFn, VarianceAccumulator testAccumulatorInput, - int expectedExtractedResult) { + Object expectedExtractedResult) { this.varianceFn = varianceFn; this.testAccumulatorInput = testAccumulatorInput;