diff --git a/sdks/java/extensions/sql/src/main/codegen/config.fmpp b/sdks/java/extensions/sql/src/main/codegen/config.fmpp index c3692430610f..7398f56b7d12 100644 --- a/sdks/java/extensions/sql/src/main/codegen/config.fmpp +++ b/sdks/java/extensions/sql/src/main/codegen/config.fmpp @@ -451,12 +451,23 @@ data: { # Return type of method implementation should be "SqlTypeNameSpec". # Example: SqlParseTimeStampZ(). dataTypeParserMethods: [ + # Spark angle-bracket collection/struct spellings (array<..>, map<..,..>, struct) so + # they parse anywhere Calcite accepts a DataType (e.g. CAST(x AS array)), not only on the + # '::' RHS. Returns SqlTypeNameSpec; emitted first in the core TypeName() choice under + # LOOKAHEAD(2). Each branch is decidable in <=2 tokens (ARRAY / MAP / IDENTIFIER-then-'<'), so + # a bare core type or a user-defined-type identifier never enters and falls through to core. + "SparkAngleType()" ] # List of methods for parsing builtin function calls. # Return type of method implementation should be "SqlNode". # Example: DateFunctionCall(). builtinFunctionCallMethods: [ + # Spark numeric type-constructor functions float(x)/double(x). + # FLOAT/DOUBLE are RESERVED keyword tokens, so they can never reach NamedFunctionCall and + # cannot be registered as UDFs (unlike DATE()/TIMESTAMP()) -- a grammar production is the only + # native route. Emits a standard CAST, which RexImpTable lowers natively. + "SparkTypeConstructorCall()" ] # List of methods for parsing extensions to "ALTER " calls. @@ -491,11 +502,17 @@ data: { ] # Binary operators tokens + # Spark/PostgreSQL infix cast token. Injected into the OPERATORS block across all lexer + # states. Maximal-munch guarantees "::" beats single ":" (COLON); the grammar's only COLON + # use is a single-colon k/v separator, so nothing is stolen. binaryOperatorsTokens: [ + "< DOUBLE_COLON: \"::\" >" ] # Binary operators initialization + # InfixCast(list, exprContext, s): emits SqlLibraryOperators.INFIX_CAST for "expr :: TYPE". extraBinaryExpressions: [ + "InfixCast" ] # List of files in @includes directory that have parser method diff --git a/sdks/java/extensions/sql/src/main/codegen/includes/parserImpls.ftl b/sdks/java/extensions/sql/src/main/codegen/includes/parserImpls.ftl index 94c0161c492c..7bd63811bbe7 100644 --- a/sdks/java/extensions/sql/src/main/codegen/includes/parserImpls.ftl +++ b/sdks/java/extensions/sql/src/main/codegen/includes/parserImpls.ftl @@ -899,4 +899,260 @@ SqlSetOptionBeam SqlSetOptionBeam(Span s, String scope) : ) } +/* + * =========================================================================================== + * Spark/PostgreSQL infix cast: expr :: TYPE + * =========================================================================================== + * + * LAYER A (clean, self-contained): the operator hookup. Wired in via the `extraBinaryExpressions` + * FreeMarker hook, which the vendored Parser.jj invokes inside Expression2's operator loop as + * `InfixCast(list, exprContext, s)` (Parser.jj:3827-3830). We push the vendored + * SqlLibraryOperators.INFIX_CAST operator (a SqlBinaryOperator named "::", kind CAST, prec 94, + * left-assoc) plus the RHS type spec onto the flat precedence list; SqlParserUtil.toTree reduces it + * to INFIX_CAST(expr, typeSpec), which is exactly how CAST(expr AS type) is represented internally. + * Precedence 94 binds tighter than arithmetic/comparison, so 1+2::int = 1+(2::int) and a::int<5 = + * (a::int)<5; chaining x::int::string = (x::int)::string falls out for free. + * + * LAYER B (SparkDataType & friends below): the Spark-aware RHS type grammar. This re-expresses + * SparkSqlPreprocessor.translateSparkType in JavaCC. *** SHARED BLAST RADIUS *** with + * rewriteAngleBracketTypes (owned by a different agent): both need the same array<>/map<>/struct<> + * + scalar-alias grammar. This copy is the merge seed; reconcile into ONE shared production before + * deleting either regex. Do NOT delete translateSparkType (still used by the angle-bracket rewrite). + */ +void InfixCast(List list, ExprContext exprContext, Span s) : +{ + final SqlDataTypeSpec dt; +} +{ + { + checkNonQueryExpression(exprContext); + } + dt = SparkDataType() { + list.add( + new SqlParserUtil.ToTreeListItem( + SqlLibraryOperators.INFIX_CAST, getPos())); + list.add(dt); + } +} + +/** + * Spark-aware data type for the RHS of "::". Tries Spark-only spellings first, then falls back to + * core DataType()'s TypeName(); keeps the SQL-standard postfix collection loop (e.g. INT ARRAY). + */ +SqlDataTypeSpec SparkDataType() : +{ + SqlTypeNameSpec tn; + final Span s = Span.of(); +} +{ + tn = SparkTypeName(s) + ( + tn = CollectionsTypeName(tn) + )* + { + return new SqlDataTypeSpec(tn, s.add(tn.getParserPos()).pos()); + } +} + +SqlTypeNameSpec SparkTypeName(Span s) : +{ + SqlTypeNameSpec t; + final SqlDataTypeSpec elem; + final SqlDataTypeSpec key; + final SqlDataTypeSpec val; + final SqlTypeName scalar; +} +{ + ( + // Spark angle-bracket ARRAY (core ARRAY is postfix). ARRAY is a case-insensitive keyword + // token; nested array> closes via two GT tokens (no ">>" shift token exists). + elem = SparkDataType() { + t = new SqlCollectionTypeNameSpec( + elem.getTypeNameSpec(), SqlTypeName.ARRAY, getPos()); + } + | + // Spark MAP with Spark-aware inner types (core MapTypeName recurses into core + // DataType, which rejects e.g. map). LA(2) so a bare MAP without '<' falls + // through to the core handler. + LOOKAHEAD(2) + key = SparkDataType() val = SparkDataType() { + t = new SqlMapTypeNameSpec(key, val, getPos()); + } + | + // Spark STRUCT. STRUCT is NOT a keyword token (it lexes as IDENTIFIER), so guard + // by image + a following '<'. + LOOKAHEAD({ getToken(1).kind == IDENTIFIER + && "STRUCT".equalsIgnoreCase(getToken(1).image) + && getToken(2).kind == LT }) + t = SparkStructType() + | + // Spark scalar aliases core DataType() rejects (string/long/short/byte/bool/ + // timestamp_ltz/timestamp_ntz), matched by IDENTIFIER image. The check is inlined as a plain + // boolean (not a JAVACODE helper) so it can run inside the generated jj_3R_* lookahead + // methods, which are not declared to throw ParseException. + LOOKAHEAD({ getToken(1).kind == IDENTIFIER && getToken(1).image != null + && (getToken(1).image.equalsIgnoreCase("STRING") + || getToken(1).image.equalsIgnoreCase("LONG") + || getToken(1).image.equalsIgnoreCase("SHORT") + || getToken(1).image.equalsIgnoreCase("BYTE") + || getToken(1).image.equalsIgnoreCase("BOOL") + || getToken(1).image.equalsIgnoreCase("TIMESTAMP_LTZ") + || getToken(1).image.equalsIgnoreCase("TIMESTAMP_NTZ")) }) + scalar = SparkScalarAlias(s) { + t = new SqlBasicTypeNameSpec(scalar, s.end(this)); + } + | + // Everything core already handles: int, bigint, decimal(p,s), varchar, char, date, + // timestamp, real, binary, MAP<>, ROW(), user-defined type, etc. + t = TypeName() + ) + { + return t; + } +} + +SqlTypeNameSpec SparkStructType() : +{ + final List names = new ArrayList(); + final List types = new ArrayList(); + SqlIdentifier n; + SqlDataTypeSpec ft; +} +{ + // "STRUCT" (guaranteed IDENTIFIER by the LOOKAHEAD at the call site). + + + n = SimpleIdentifier() ft = SparkDataType() { + names.add(n); + types.add(ft); + } + ( + n = SimpleIdentifier() ft = SparkDataType() { + names.add(n); + types.add(ft); + } + )* + + { + return new SqlRowTypeNameSpec(getPos(), names, types); + } +} + +SqlTypeName SparkScalarAlias(Span s) : +{ + final SqlTypeName tn; +} +{ + // Only reached when isSparkScalarAlias(getToken(1)) already matched, so the switch is total. + { + s.add(this); + switch (token.image.toUpperCase(Locale.ROOT)) { + case "STRING": tn = SqlTypeName.VARCHAR; break; + case "LONG": tn = SqlTypeName.BIGINT; break; + case "SHORT": tn = SqlTypeName.SMALLINT; break; + case "BYTE": tn = SqlTypeName.TINYINT; break; + case "BOOL": tn = SqlTypeName.BOOLEAN; break; + case "TIMESTAMP_LTZ": + case "TIMESTAMP_NTZ": tn = SqlTypeName.TIMESTAMP; break; + default: + throw new ParseException("not a Spark scalar type alias: " + token.image); + } + return tn; + } +} + +/** + * Spark numeric type-constructor functions float(x) / double(x). + * + *

FLOAT and DOUBLE are RESERVED keyword tokens, so the lexer emits / (never + * IDENTIFIER) and the call can never reach NamedFunctionCall / function resolution -- registering + * them as scalar UDFs is therefore impossible (the route DATE()/TIMESTAMP() take). A dedicated + * builtin-function production is the only native option. We emit a standard CAST(expr AS FLOAT/ + * DOUBLE), which RexImpTable lowers natively. Strictly better than the regex it retires, whose + * non-nesting [^)]+? group broke on nested parens (float(abs(x))). + * + *

Hooked in via the `builtinFunctionCallMethods` FreeMarker list, so this is one alternative of + * the generated BuiltinFunctionCall(); / are unique first tokens there, so there is + * no choice conflict. + */ +SqlNode SparkTypeConstructorCall() : +{ + final Span s; + final SqlNode e; + final SqlTypeName tn; +} +{ + ( { tn = SqlTypeName.FLOAT; } | { tn = SqlTypeName.DOUBLE; } ) + { s = span(); } + e = Expression(ExprContext.ACCEPT_SUB_QUERY) + { + return SqlStdOperatorTable.CAST.createCall( + s.end(this), + e, + new SqlDataTypeSpec(new SqlBasicTypeNameSpec(tn, getPos()), getPos())); + } +} + +/** + * Spark angle-bracket collection / struct types in ANY DataType position: array, map, + * struct. Returns SqlTypeNameSpec so it can be wired into the core TypeName() choice via the + * dataTypeParserMethods hook (config.fmpp), making these spellings parse in CAST(x AS array) + * etc. -- not only on the '::' RHS (where SparkTypeName already accepts them). Inner element / key / + * value / field types use SparkDataType() so Spark scalar spellings resolve inside the brackets + * (array, map). + * + *

Each branch is decidable in <=2 tokens (ARRAY / MAP keyword, or IDENTIFIER-then-'<' for struct), + * so under the template's LOOKAHEAD(2) wrapper a bare core type (a keyword token) or a user-defined + * type identifier (IDENTIFIER not followed by '<') never enters here and falls through to core + * SqlTypeName() / CompoundIdentifier(). No semantic guard on the struct branch: struct is the only + * IDENTIFIER-initial alternative, so IDENTIFIER-'<' in a type position is unambiguously a struct. + */ +SqlTypeNameSpec SparkAngleType() : +{ + SqlTypeNameSpec t; + final SqlDataTypeSpec elem; + final SqlDataTypeSpec key; + final SqlDataTypeSpec val; + final SqlTypeName scalar; + final Span s = Span.of(); +} +{ + ( + elem = SparkDataType() { + t = new SqlCollectionTypeNameSpec( + elem.getTypeNameSpec(), SqlTypeName.ARRAY, getPos()); + } + | + key = SparkDataType() val = SparkDataType() { + t = new SqlMapTypeNameSpec(key, val, getPos()); + } + | + // Spark scalar aliases (string/long/short/byte/bool/timestamp_ltz/timestamp_ntz) in ANY + // DataType position -- e.g. CAST(x AS STRING) -- not only on the '::' RHS (where + // SparkTypeName already accepts them). SparkScalarAlias maps STRING->VARCHAR, LONG->BIGINT, + // etc. (line ~1041). This is the native, nested-paren-safe replacement for the regex + // CAST(.. AS STRING)->CAST(.. AS VARCHAR) rewrite that lived in SparkSqlPreprocessor. + // Guarded by an inlined boolean LOOKAHEAD (mirroring SparkTypeName's branch) so it runs + // inside the generated jj_3R_* lookahead methods, which are not declared to throw + // ParseException; the wrapping core TypeName() LOOKAHEAD(2) only commits here when token 1 + // is one of these aliases (an IDENTIFIER, never a keyword core type). + LOOKAHEAD({ getToken(1).kind == IDENTIFIER && getToken(1).image != null + && (getToken(1).image.equalsIgnoreCase("STRING") + || getToken(1).image.equalsIgnoreCase("LONG") + || getToken(1).image.equalsIgnoreCase("SHORT") + || getToken(1).image.equalsIgnoreCase("BYTE") + || getToken(1).image.equalsIgnoreCase("BOOL") + || getToken(1).image.equalsIgnoreCase("TIMESTAMP_LTZ") + || getToken(1).image.equalsIgnoreCase("TIMESTAMP_NTZ")) }) + scalar = SparkScalarAlias(s) { + t = new SqlBasicTypeNameSpec(scalar, s.end(this)); + } + | + t = SparkStructType() + ) + { + return t; + } +} + // End parserImpls.ftl diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/CalciteQueryPlanner.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/CalciteQueryPlanner.java index 606a3c5f71a2..81c62ad42964 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/CalciteQueryPlanner.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/CalciteQueryPlanner.java @@ -153,7 +153,25 @@ public FrameworkConfig defaultConfig(JdbcConnection connection, Collection getRuleSets() { + return ImmutableList.of(RuleSets.ofList(getAllRules())); + } - return ImmutableList.of( - RuleSets.ofList( - ImmutableList.builder() - .addAll(BEAM_CONVERTERS) - .addAll(BEAM_TO_ENUMERABLE) - .addAll(LOGICAL_OPTIMIZATIONS) - .build())); + public static List getAllRules() { + return ImmutableList.builder() + .addAll(BEAM_CONVERTERS) + .addAll(BEAM_TO_ENUMERABLE) + .addAll(LOGICAL_OPTIMIZATIONS) + .build(); } } diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamWindowRel.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamWindowRel.java index e597b6c63ed4..2589b52e22e4 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamWindowRel.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamWindowRel.java @@ -142,9 +142,23 @@ public PTransform, PCollection> buildPTransform() { List argList = anAggCall.getArgList(); Schema.Field field = CalciteUtils.toField(anAggCall.getName(), anAggCall.getType()); - Combine.CombineFn combineFn = - AggregationCombineFnAdapter.createCombineFnAnalyticsFunctions( - anAggCall, field, anAggCall.getAggregation().getName()); + String aggName = anAggCall.getAggregation().getName(); + Combine.CombineFn combineFn; + if (anAggCall.ignoreNulls() + && ("FIRST_VALUE".equals(aggName) || "LAST_VALUE".equals(aggName))) { + // Spark's first(col, ignoreNulls=true) / last(col, ignoreNulls=true) + // (as emitted by pandas ffill / bfill) must skip nulls within the frame. + combineFn = + "FIRST_VALUE".equals(aggName) + ? org.apache.beam.sdk.extensions.sql.impl.transform + .BeamBuiltinAnalyticFunctions.navigationFirstValue(true) + : org.apache.beam.sdk.extensions.sql.impl.transform + .BeamBuiltinAnalyticFunctions.navigationLastValue(true); + } else { + combineFn = + AggregationCombineFnAdapter.createCombineFnAnalyticsFunctions( + anAggCall, field, aggName); + } FieldAggregation fieldAggregation = new FieldAggregation( partitionKeysDef, @@ -208,8 +222,12 @@ public FieldAggregation( @Override public NodeStats estimateNodeStats(BeamRelMetadataQuery mq) { - NodeStats inputStat = BeamSqlRelUtils.getNodeStats(this.input, mq); - return inputStat; + // Derive from the planner's row-count for the input, which stays finite even when the input is + // an unoptimized RelSubset (BeamSqlRelUtils.getNodeStats returns UNKNOWN for a RelSubset, which + // makes the cumulative cost infinite and aborts planning -- notably for an OVER nested inside a + // projection expression, as emitted by pandas ffill / bfill). + double rows = mq.getRowCount(this.input); + return NodeStats.create(rows, 0d, rows); } /** @@ -220,10 +238,12 @@ public NodeStats estimateNodeStats(BeamRelMetadataQuery mq) { */ @Override public BeamCostModel beamComputeSelfCost(RelOptPlanner planner, BeamRelMetadataQuery mq) { - NodeStats inputStat = BeamSqlRelUtils.getNodeStats(this.input, mq); + // Use the planner's row-count (RelSubset-safe) rather than BeamSqlRelUtils.getNodeStats, which + // yields UNKNOWN/infinite for a RelSubset input and prevents the Volcano cost fixpoint from + // converging. + double rows = mq.getRowCount(this.input); float multiplier = 1f + 0.125f; - return BeamCostModel.FACTORY.makeCost( - inputStat.getRowCount() * multiplier, inputStat.getRate() * multiplier); + return BeamCostModel.FACTORY.makeCost(rows * multiplier, 0d); } private static class Transform extends PTransform, PCollection> { @@ -317,8 +337,12 @@ public void processElement( (long) idx, (long) sortedRowsAsList.size()); } else { - accumulator = - fieldAgg.combineFn.addInput(accumulator, aggRow.getBaseValue(aggFieldIndex)); + // Functions with no input column (e.g. COUNT(*) OVER (...)) carry no field index; + // feed the combiner the (always non-null) row itself so COUNT counts every row, + // rather than indexing the row with -1. + Object inputValue = + aggFieldIndex >= 0 ? aggRow.getBaseValue(aggFieldIndex) : (Object) aggRow; + accumulator = fieldAgg.combineFn.addInput(accumulator, inputValue); } count++; } @@ -395,8 +419,12 @@ private List getRows(List input, int index) { fieldAgg.upperLimit != null ? fieldAgg.upperLimit.intValue() : Integer.MAX_VALUE; int lowerIndex = ll == Integer.MAX_VALUE ? Integer.MIN_VALUE : index - ll; int upperIndex = ul == Integer.MAX_VALUE ? Integer.MAX_VALUE : index + ul + 1; - lowerIndex = lowerIndex < 0 ? 0 : lowerIndex; - upperIndex = upperIndex > input.size() ? input.size() : upperIndex; + // Clamp both bounds into [0, size]. A frame that lies entirely before the partition + // (upperIndex < 0, e.g. ROWS BETWEEN 2 PRECEDING AND 2 PRECEDING on the first rows, as + // emitted by pandas pct_change/diff with periods>1) or entirely after it (lowerIndex > + // size, the FOLLOWING analogue) must yield an empty frame rather than throw on subList. + lowerIndex = Math.max(0, Math.min(lowerIndex, input.size())); + upperIndex = Math.max(lowerIndex, Math.min(upperIndex, input.size())); List out = input.subList(lowerIndex, upperIndex); return out; } diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rule/BeamWindowRule.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rule/BeamWindowRule.java index c442ae3d1947..ee178b9ebaf2 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rule/BeamWindowRule.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rule/BeamWindowRule.java @@ -36,10 +36,16 @@ private BeamWindowRule() { @Override public RelNode convert(RelNode relNode) { Window w = (Window) relNode; + RelNode input = w.getInput(); return new BeamWindowRel( w.getCluster(), w.getTraitSet().replace(BeamLogicalConvention.INSTANCE), - w.getInput(), + // Request a BEAM_LOGICAL-convention input rather than handing the BeamWindowRel the raw + // NONE-convention RelSubset. Without this, when the window is reached via CALC_TO_WINDOW + // (its input is a fresh Calc/Project), the input subset is never converted to BEAM_LOGICAL, + // its best cost stays infinite, and planning fails. Mirrors BeamSortRule / + // BeamAggregationRule. + convert(input, input.getTraitSet().replace(BeamLogicalConvention.INSTANCE)), w.getConstants(), w.getRowType(), w.groups); diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/schema/BeamTableUtils.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/schema/BeamTableUtils.java index d9dbf881850e..68cfe7f3a612 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/schema/BeamTableUtils.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/schema/BeamTableUtils.java @@ -22,6 +22,8 @@ import java.io.IOException; import java.io.StringWriter; import java.math.BigDecimal; +import java.math.RoundingMode; +import java.time.Duration; import java.time.Instant; import java.time.LocalDate; import java.time.LocalTime; @@ -34,6 +36,7 @@ import org.apache.beam.sdk.schemas.Schema; import org.apache.beam.sdk.schemas.Schema.FieldType; import org.apache.beam.sdk.schemas.Schema.TypeName; +import org.apache.beam.sdk.schemas.logicaltypes.NanosDuration; import org.apache.beam.sdk.values.Row; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.avatica.util.ByteString; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.util.NlsString; @@ -146,6 +149,15 @@ public static Object autoCastField(Schema.Field field, @Nullable Object rawObj) } else if (CalciteUtils.isDateTimeType(type)) { // Internal representation of Date in Calcite is convertible to Joda's Datetime. return new DateTime(rawObj); + } else if (type.getTypeName().isLogicalType() + && type.getLogicalType() != null + && NanosDuration.IDENTIFIER.equals(type.getLogicalType().getIdentifier()) + && rawObj instanceof BigDecimal) { + // Calcite carries an INTERVAL DAY TO SECOND value as a (possibly fractional) millisecond + // BigDecimal; decode it back to the java.time.Duration the NanosDuration logical type + // expects, preserving sub-millisecond precision. + BigDecimal nanos = ((BigDecimal) rawObj).movePointRight(6).setScale(0, RoundingMode.HALF_UP); + return Duration.ofNanos(nanos.longValueExact()); // handle decimal } else if (type.getTypeName().isNumericType() && ((rawObj instanceof String) diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/transform/BeamBuiltinAnalyticFunctions.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/transform/BeamBuiltinAnalyticFunctions.java index 413093d416d8..2596c5fdfa96 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/transform/BeamBuiltinAnalyticFunctions.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/transform/BeamBuiltinAnalyticFunctions.java @@ -45,6 +45,7 @@ public class BeamBuiltinAnalyticFunctions { .put("DENSE_RANK", typeName -> numberingDenseRank()) .put("RANK", typeName -> numberingRank()) .put("PERCENT_RANK", typeName -> numberingPercentRank()) + .put("CUME_DIST", typeName -> numberingCumeDist()) .build(); public static Combine.CombineFn create(String functionName, Schema.FieldType fieldType) { @@ -66,9 +67,97 @@ public class BeamBuiltinAnalyticFunctions { return new LastValueCombineFn(); } - private static class FirstValueCombineFn extends Combine.CombineFn, T> { + /** + * FIRST_VALUE that optionally skips nulls. With {@code ignoreNulls=true} it returns the first + * non-null value in the frame (matching Spark's {@code first(col, ignoreNulls=true)}, as emitted + * by pandas ffill / bfill); otherwise it behaves like {@link #navigationFirstValue()}. + */ + public static Combine.CombineFn navigationFirstValue(boolean ignoreNulls) { + return ignoreNulls ? new FirstValueIgnoreNullsCombineFn() : new FirstValueCombineFn(); + } + + /** + * LAST_VALUE that optionally skips nulls. With {@code ignoreNulls=true} it returns the last + * non-null value in the frame (matching Spark's {@code last(col, ignoreNulls=true)}); otherwise + * it behaves like {@link #navigationLastValue()}. + */ + public static Combine.CombineFn navigationLastValue(boolean ignoreNulls) { + return ignoreNulls ? new LastValueIgnoreNullsCombineFn() : new LastValueCombineFn(); + } + + /** + * Accumulator for FIRST_VALUE / LAST_VALUE that, unlike {@link Optional}, can represent a {@code + * null} value distinct from "no row seen yet". Plain {@code Optional.of(input)} throws an NPE + * when a row's value is null, but FIRST_VALUE / LAST_VALUE (without IGNORE NULLS) must faithfully + * return the first / last row's value even when that value is null. + */ + private static class ValueHolder { + private boolean seen; + private T value; + + void set(T input) { + this.seen = true; + this.value = input; + } + } + + private static class FirstValueCombineFn extends Combine.CombineFn, T> { private FirstValueCombineFn() {} + @Override + public ValueHolder createAccumulator() { + return new ValueHolder<>(); + } + + @Override + public ValueHolder addInput(ValueHolder accumulator, T input) { + if (!accumulator.seen) { + accumulator.set(input); + } + return accumulator; + } + + @Override + public ValueHolder mergeAccumulators(Iterable> accumulators) { + throw new UnsupportedOperationException(); + } + + @Override + public T extractOutput(ValueHolder accumulator) { + return accumulator.seen ? accumulator.value : null; + } + } + + private static class LastValueCombineFn extends Combine.CombineFn, T> { + private LastValueCombineFn() {} + + @Override + public ValueHolder createAccumulator() { + return new ValueHolder<>(); + } + + @Override + public ValueHolder addInput(ValueHolder accumulator, T input) { + accumulator.set(input); + return accumulator; + } + + @Override + public ValueHolder mergeAccumulators(Iterable> accumulators) { + throw new UnsupportedOperationException(); + } + + @Override + public T extractOutput(ValueHolder accumulator) { + return accumulator.seen ? accumulator.value : null; + } + } + + /** FIRST_VALUE(... IGNORE NULLS): the first non-null value in the frame. */ + private static class FirstValueIgnoreNullsCombineFn + extends Combine.CombineFn, T> { + private FirstValueIgnoreNullsCombineFn() {} + @Override public Optional createAccumulator() { return Optional.empty(); @@ -76,11 +165,10 @@ public Optional createAccumulator() { @Override public Optional addInput(Optional accumulator, T input) { - Optional r = accumulator; - if (!accumulator.isPresent()) { - r = Optional.of(input); + if (!accumulator.isPresent() && input != null) { + return Optional.of(input); } - return r; + return accumulator; } @Override @@ -94,8 +182,10 @@ public T extractOutput(Optional accumulator) { } } - private static class LastValueCombineFn extends Combine.CombineFn, T> { - private LastValueCombineFn() {} + /** LAST_VALUE(... IGNORE NULLS): the last non-null value in the frame. */ + private static class LastValueIgnoreNullsCombineFn + extends Combine.CombineFn, T> { + private LastValueIgnoreNullsCombineFn() {} @Override public Optional createAccumulator() { @@ -104,8 +194,10 @@ public Optional createAccumulator() { @Override public Optional addInput(Optional accumulator, T input) { - Optional r = Optional.of(input); - return r; + if (input != null) { + return Optional.of(input); + } + return accumulator; } @Override @@ -136,6 +228,10 @@ public T extractOutput(Optional accumulator) { return new PercentRankCombineFn(); } + public static Combine.CombineFn numberingCumeDist() { + return new CumeDistCombineFn(); + } + public abstract static class PositionAwareCombineFn extends Combine.CombineFn { public abstract AccumT addInput( @@ -251,6 +347,45 @@ public Long extractOutput(KV accumulator) { } } + /** + * CUME_DIST: the relative rank of the current row, computed as {@code (number of rows preceding + * or peer with the current row) / (total rows in the partition)}. + * + *

Beam evaluates analytic functions over the per-row frame. With the SQL-standard default + * frame for CUME_DIST (RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW), the frame handed to + * this combiner already contains exactly the "preceding or peer" rows, so the numerator is simply + * the number of {@code addInput} calls. The denominator is the partition size, which is supplied + * to every {@code addInput} call as {@code countPartition}. + */ + private static class CumeDistCombineFn + extends PositionAwareCombineFn, Double> { + + @Override + public KV createAccumulator() { + // KV(rowsInFrameSoFar, partitionSize) + return KV.of(0L, 0L); + } + + @Override + public KV addInput( + KV accumulator, + BigDecimal input, + Long cursorPosition, + Long cursorPartition, + Long countPartition) { + return KV.of(accumulator.getKey() + 1L, countPartition); + } + + @Override + public Double extractOutput(KV accumulator) { + long total = accumulator.getValue(); + if (total <= 0L) { + return 0.0; + } + return accumulator.getKey().doubleValue() / (double) total; + } + } + private static class PercentRankCombineFn extends PositionAwareCombineFn, KV>, Double> { diff --git a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamAnalyticFunctionsTest.java b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamAnalyticFunctionsTest.java index d66fdb58d4af..5ca65623c32b 100644 --- a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamAnalyticFunctionsTest.java +++ b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamAnalyticFunctionsTest.java @@ -588,4 +588,58 @@ public void testPercentRankFunction() throws Exception { pipeline.run(); } + + @Test + public void testCumeDistFunction() throws Exception { + pipeline.enableAbandonedNodeEnforcement(false); + PCollection inputRows = inputData2(); + String sql = "SELECT x, CUME_DIST() over (ORDER BY x ) as agg FROM PCOLLECTION"; + PCollection result = inputRows.apply("sql", SqlTransform.query(sql)); + + Schema overResultSchema = Schema.builder().addInt32Field("x").addDoubleField("agg").build(); + + // CUME_DIST = (# rows with value <= current) / (total rows). Input x: 1,2,2,5,8,10,10. + List overResult = + TestUtils.RowsBuilder.of(overResultSchema) + .addRows( + 1, 1.0 / 7.0, + 2, 3.0 / 7.0, + 2, 3.0 / 7.0, + 5, 4.0 / 7.0, + 8, 5.0 / 7.0, + 10, 7.0 / 7.0, + 10, 7.0 / 7.0) + .getRows(); + + PAssert.that(result).containsInAnyOrder(overResult); + + pipeline.run(); + } + + @Test + public void testCountStarOverWindow() throws Exception { + pipeline.enableAbandonedNodeEnforcement(false); + PCollection inputRows = inputData2(); + String sql = "SELECT x, COUNT(*) over () as agg FROM PCOLLECTION"; + PCollection result = inputRows.apply("sql", SqlTransform.query(sql)); + + Schema overResultSchema = Schema.builder().addInt32Field("x").addInt64Field("agg").build(); + + // COUNT(*) OVER () counts every row in the partition (7 rows) regardless of value. + List overResult = + TestUtils.RowsBuilder.of(overResultSchema) + .addRows( + 1, 7L, + 2, 7L, + 2, 7L, + 5, 7L, + 8, 7L, + 10, 7L, + 10, 7L) + .getRows(); + + PAssert.that(result).containsInAnyOrder(overResult); + + pipeline.run(); + } } diff --git a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlDslArrayTest.java b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlDslArrayTest.java index 65d6d72d657c..3bac626ac3f0 100644 --- a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlDslArrayTest.java +++ b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlDslArrayTest.java @@ -68,6 +68,29 @@ public void testSelectArrayValue() { pipeline.run(); } + @Test + public void testSelectArrayFunctionForm() { + PCollection input = pCollectionOf2Elements(); + + Schema resultType = + Schema.builder() + .addInt32Field("f_int") + .addArrayField("f_arr", Schema.FieldType.STRING) + .build(); + + PCollection result = + input.apply( + "sqlQuery", + SqlTransform.query("SELECT 42, ARRAY('aa', 'bb') as `f_arr` FROM PCOLLECTION")); + + PAssert.that(result) + .containsInAnyOrder( + Row.withSchema(resultType).addValues(42, Arrays.asList("aa", "bb")).build(), + Row.withSchema(resultType).addValues(42, Arrays.asList("aa", "bb")).build()); + + pipeline.run(); + } + @Test public void testProjectArrayField() { PCollection input = pCollectionOf2Elements(); diff --git a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlMapTest.java b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlMapTest.java index 4a2ad664b4b3..09e75d67aa2e 100644 --- a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlMapTest.java +++ b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlMapTest.java @@ -90,6 +90,28 @@ public void testSelectMapField() { pipeline.run(); } + @Test + public void testSelectMapFunctionForm() { + PCollection input = pCollectionOf2Elements(); + + Schema resultType = + Schema.builder() + .addInt32Field("f_int") + .addMapField("f_map", Schema.FieldType.STRING, Schema.FieldType.INT32) + .build(); + + PCollection result = + input.apply( + "sqlQuery", SqlTransform.query("SELECT 42, MAP('aa', 1) as `f_map` FROM PCOLLECTION")); + + PAssert.that(result) + .containsInAnyOrder( + Row.withSchema(resultType).addValues(42, ImmutableMap.of("aa", 1)).build(), + Row.withSchema(resultType).addValues(42, ImmutableMap.of("aa", 1)).build()); + + pipeline.run(); + } + @Test public void testSelectMapFieldKeyValueSameType() { PCollection input = pCollectionOf2Elements();