Skip to content

Commit 211ee94

Browse files
damccormdamccorm
committed
Implement Spark SQL parser extensions and enhance analytic functions
1 parent 9bcdedd commit 211ee94

8 files changed

Lines changed: 544 additions & 29 deletions

File tree

‎sdks/java/extensions/sql/src/main/codegen/config.fmpp‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -451,12 +451,23 @@ data: {
451451
# Return type of method implementation should be "SqlTypeNameSpec".
452452
# Example: SqlParseTimeStampZ().
453453
dataTypeParserMethods: [
454+
# Spark angle-bracket collection/struct spellings (array<..>, map<..,..>, struct<f:T,..>) so
455+
# they parse anywhere Calcite accepts a DataType (e.g. CAST(x AS array<int>)), not only on the
456+
# '::' RHS. Returns SqlTypeNameSpec; emitted first in the core TypeName() choice under
457+
# LOOKAHEAD(2). Each branch is decidable in <=2 tokens (ARRAY / MAP / IDENTIFIER-then-'<'), so
458+
# a bare core type or a user-defined-type identifier never enters and falls through to core.
459+
"SparkAngleType()"
454460
]
455461

456462
# List of methods for parsing builtin function calls.
457463
# Return type of method implementation should be "SqlNode".
458464
# Example: DateFunctionCall().
459465
builtinFunctionCallMethods: [
466+
# Spark numeric type-constructor functions float(x)/double(x).
467+
# FLOAT/DOUBLE are RESERVED keyword tokens, so they can never reach NamedFunctionCall and
468+
# cannot be registered as UDFs (unlike DATE()/TIMESTAMP()) -- a grammar production is the only
469+
# native route. Emits a standard CAST, which RexImpTable lowers natively.
470+
"SparkTypeConstructorCall()"
460471
]
461472

462473
# List of methods for parsing extensions to "ALTER <scope>" calls.
@@ -491,11 +502,17 @@ data: {
491502
]
492503

493504
# Binary operators tokens
505+
# Spark/PostgreSQL infix cast token. Injected into the OPERATORS block across all lexer
506+
# states. Maximal-munch guarantees "::" beats single ":" (COLON); the grammar's only COLON
507+
# use is a single-colon k/v separator, so nothing is stolen.
494508
binaryOperatorsTokens: [
509+
"< DOUBLE_COLON: \"::\" >"
495510
]
496511

497512
# Binary operators initialization
513+
# InfixCast(list, exprContext, s): emits SqlLibraryOperators.INFIX_CAST for "expr :: TYPE".
498514
extraBinaryExpressions: [
515+
"InfixCast"
499516
]
500517

501518
# List of files in @includes directory that have parser method

‎sdks/java/extensions/sql/src/main/codegen/includes/parserImpls.ftl‎

Lines changed: 256 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -899,4 +899,260 @@ SqlSetOptionBeam SqlSetOptionBeam(Span s, String scope) :
899899
)
900900
}
901901

902+
/*
903+
* ===========================================================================================
904+
* Spark/PostgreSQL infix cast: expr :: TYPE
905+
* ===========================================================================================
906+
*
907+
* LAYER A (clean, self-contained): the operator hookup. Wired in via the `extraBinaryExpressions`
908+
* FreeMarker hook, which the vendored Parser.jj invokes inside Expression2's operator loop as
909+
* `InfixCast(list, exprContext, s)` (Parser.jj:3827-3830). We push the vendored
910+
* SqlLibraryOperators.INFIX_CAST operator (a SqlBinaryOperator named "::", kind CAST, prec 94,
911+
* left-assoc) plus the RHS type spec onto the flat precedence list; SqlParserUtil.toTree reduces it
912+
* to INFIX_CAST(expr, typeSpec), which is exactly how CAST(expr AS type) is represented internally.
913+
* Precedence 94 binds tighter than arithmetic/comparison, so 1+2::int = 1+(2::int) and a::int<5 =
914+
* (a::int)<5; chaining x::int::string = (x::int)::string falls out for free.
915+
*
916+
* LAYER B (SparkDataType & friends below): the Spark-aware RHS type grammar. This re-expresses
917+
* SparkSqlPreprocessor.translateSparkType in JavaCC. *** SHARED BLAST RADIUS *** with
918+
* rewriteAngleBracketTypes (owned by a different agent): both need the same array<>/map<>/struct<>
919+
* + scalar-alias grammar. This copy is the merge seed; reconcile into ONE shared production before
920+
* deleting either regex. Do NOT delete translateSparkType (still used by the angle-bracket rewrite).
921+
*/
922+
void InfixCast(List<Object> list, ExprContext exprContext, Span s) :
923+
{
924+
final SqlDataTypeSpec dt;
925+
}
926+
{
927+
<DOUBLE_COLON> {
928+
checkNonQueryExpression(exprContext);
929+
}
930+
dt = SparkDataType() {
931+
list.add(
932+
new SqlParserUtil.ToTreeListItem(
933+
SqlLibraryOperators.INFIX_CAST, getPos()));
934+
list.add(dt);
935+
}
936+
}
937+
938+
/**
939+
* Spark-aware data type for the RHS of "::". Tries Spark-only spellings first, then falls back to
940+
* core DataType()'s TypeName(); keeps the SQL-standard postfix collection loop (e.g. INT ARRAY).
941+
*/
942+
SqlDataTypeSpec SparkDataType() :
943+
{
944+
SqlTypeNameSpec tn;
945+
final Span s = Span.of();
946+
}
947+
{
948+
tn = SparkTypeName(s)
949+
(
950+
tn = CollectionsTypeName(tn)
951+
)*
952+
{
953+
return new SqlDataTypeSpec(tn, s.add(tn.getParserPos()).pos());
954+
}
955+
}
956+
957+
SqlTypeNameSpec SparkTypeName(Span s) :
958+
{
959+
SqlTypeNameSpec t;
960+
final SqlDataTypeSpec elem;
961+
final SqlDataTypeSpec key;
962+
final SqlDataTypeSpec val;
963+
final SqlTypeName scalar;
964+
}
965+
{
966+
(
967+
// Spark angle-bracket ARRAY (core ARRAY is postfix). ARRAY is a case-insensitive keyword
968+
// token; nested array<array<int>> closes via two GT tokens (no ">>" shift token exists).
969+
<ARRAY> <LT> elem = SparkDataType() <GT> {
970+
t = new SqlCollectionTypeNameSpec(
971+
elem.getTypeNameSpec(), SqlTypeName.ARRAY, getPos());
972+
}
973+
|
974+
// Spark MAP<K,V> with Spark-aware inner types (core MapTypeName recurses into core
975+
// DataType, which rejects e.g. map<string,int>). LA(2) so a bare MAP without '<' falls
976+
// through to the core handler.
977+
LOOKAHEAD(2)
978+
<MAP> <LT> key = SparkDataType() <COMMA> val = SparkDataType() <GT> {
979+
t = new SqlMapTypeNameSpec(key, val, getPos());
980+
}
981+
|
982+
// Spark STRUCT<f:T, ...>. STRUCT is NOT a keyword token (it lexes as IDENTIFIER), so guard
983+
// by image + a following '<'.
984+
LOOKAHEAD({ getToken(1).kind == IDENTIFIER
985+
&& "STRUCT".equalsIgnoreCase(getToken(1).image)
986+
&& getToken(2).kind == LT })
987+
t = SparkStructType()
988+
|
989+
// Spark scalar aliases core DataType() rejects (string/long/short/byte/bool/
990+
// timestamp_ltz/timestamp_ntz), matched by IDENTIFIER image. The check is inlined as a plain
991+
// boolean (not a JAVACODE helper) so it can run inside the generated jj_3R_* lookahead
992+
// methods, which are not declared to throw ParseException.
993+
LOOKAHEAD({ getToken(1).kind == IDENTIFIER && getToken(1).image != null
994+
&& (getToken(1).image.equalsIgnoreCase("STRING")
995+
|| getToken(1).image.equalsIgnoreCase("LONG")
996+
|| getToken(1).image.equalsIgnoreCase("SHORT")
997+
|| getToken(1).image.equalsIgnoreCase("BYTE")
998+
|| getToken(1).image.equalsIgnoreCase("BOOL")
999+
|| getToken(1).image.equalsIgnoreCase("TIMESTAMP_LTZ")
1000+
|| getToken(1).image.equalsIgnoreCase("TIMESTAMP_NTZ")) })
1001+
scalar = SparkScalarAlias(s) {
1002+
t = new SqlBasicTypeNameSpec(scalar, s.end(this));
1003+
}
1004+
|
1005+
// Everything core already handles: int, bigint, decimal(p,s), varchar, char, date,
1006+
// timestamp, real, binary, MAP<>, ROW(), user-defined type, etc.
1007+
t = TypeName()
1008+
)
1009+
{
1010+
return t;
1011+
}
1012+
}
1013+
1014+
SqlTypeNameSpec SparkStructType() :
1015+
{
1016+
final List<SqlIdentifier> names = new ArrayList<SqlIdentifier>();
1017+
final List<SqlDataTypeSpec> types = new ArrayList<SqlDataTypeSpec>();
1018+
SqlIdentifier n;
1019+
SqlDataTypeSpec ft;
1020+
}
1021+
{
1022+
// "STRUCT" (guaranteed IDENTIFIER by the LOOKAHEAD at the call site).
1023+
<IDENTIFIER>
1024+
<LT>
1025+
n = SimpleIdentifier() <COLON> ft = SparkDataType() {
1026+
names.add(n);
1027+
types.add(ft);
1028+
}
1029+
(
1030+
<COMMA> n = SimpleIdentifier() <COLON> ft = SparkDataType() {
1031+
names.add(n);
1032+
types.add(ft);
1033+
}
1034+
)*
1035+
<GT>
1036+
{
1037+
return new SqlRowTypeNameSpec(getPos(), names, types);
1038+
}
1039+
}
1040+
1041+
SqlTypeName SparkScalarAlias(Span s) :
1042+
{
1043+
final SqlTypeName tn;
1044+
}
1045+
{
1046+
// Only reached when isSparkScalarAlias(getToken(1)) already matched, so the switch is total.
1047+
<IDENTIFIER> {
1048+
s.add(this);
1049+
switch (token.image.toUpperCase(Locale.ROOT)) {
1050+
case "STRING": tn = SqlTypeName.VARCHAR; break;
1051+
case "LONG": tn = SqlTypeName.BIGINT; break;
1052+
case "SHORT": tn = SqlTypeName.SMALLINT; break;
1053+
case "BYTE": tn = SqlTypeName.TINYINT; break;
1054+
case "BOOL": tn = SqlTypeName.BOOLEAN; break;
1055+
case "TIMESTAMP_LTZ":
1056+
case "TIMESTAMP_NTZ": tn = SqlTypeName.TIMESTAMP; break;
1057+
default:
1058+
throw new ParseException("not a Spark scalar type alias: " + token.image);
1059+
}
1060+
return tn;
1061+
}
1062+
}
1063+
1064+
/**
1065+
* Spark numeric type-constructor functions float(x) / double(x).
1066+
*
1067+
* <p>FLOAT and DOUBLE are RESERVED keyword tokens, so the lexer emits <FLOAT>/<DOUBLE> (never
1068+
* IDENTIFIER) and the call can never reach NamedFunctionCall / function resolution -- registering
1069+
* them as scalar UDFs is therefore impossible (the route DATE()/TIMESTAMP() take). A dedicated
1070+
* builtin-function production is the only native option. We emit a standard CAST(expr AS FLOAT/
1071+
* DOUBLE), which RexImpTable lowers natively. Strictly better than the regex it retires, whose
1072+
* non-nesting [^)]+? group broke on nested parens (float(abs(x))).
1073+
*
1074+
* <p>Hooked in via the `builtinFunctionCallMethods` FreeMarker list, so this is one alternative of
1075+
* the generated BuiltinFunctionCall(); <FLOAT>/<DOUBLE> are unique first tokens there, so there is
1076+
* no choice conflict.
1077+
*/
1078+
SqlNode SparkTypeConstructorCall() :
1079+
{
1080+
final Span s;
1081+
final SqlNode e;
1082+
final SqlTypeName tn;
1083+
}
1084+
{
1085+
( <FLOAT> { tn = SqlTypeName.FLOAT; } | <DOUBLE> { tn = SqlTypeName.DOUBLE; } )
1086+
{ s = span(); }
1087+
<LPAREN> e = Expression(ExprContext.ACCEPT_SUB_QUERY) <RPAREN>
1088+
{
1089+
return SqlStdOperatorTable.CAST.createCall(
1090+
s.end(this),
1091+
e,
1092+
new SqlDataTypeSpec(new SqlBasicTypeNameSpec(tn, getPos()), getPos()));
1093+
}
1094+
}
1095+
1096+
/**
1097+
* Spark angle-bracket collection / struct types in ANY DataType position: array<E>, map<K,V>,
1098+
* struct<f:T,..>. Returns SqlTypeNameSpec so it can be wired into the core TypeName() choice via the
1099+
* dataTypeParserMethods hook (config.fmpp), making these spellings parse in CAST(x AS array<int>)
1100+
* etc. -- not only on the '::' RHS (where SparkTypeName already accepts them). Inner element / key /
1101+
* value / field types use SparkDataType() so Spark scalar spellings resolve inside the brackets
1102+
* (array<string>, map<string,int>).
1103+
*
1104+
* <p>Each branch is decidable in <=2 tokens (ARRAY / MAP keyword, or IDENTIFIER-then-'<' for struct),
1105+
* so under the template's LOOKAHEAD(2) wrapper a bare core type (a keyword token) or a user-defined
1106+
* type identifier (IDENTIFIER not followed by '<') never enters here and falls through to core
1107+
* SqlTypeName() / CompoundIdentifier(). No semantic guard on the struct branch: struct is the only
1108+
* IDENTIFIER-initial alternative, so IDENTIFIER-'<' in a type position is unambiguously a struct.
1109+
*/
1110+
SqlTypeNameSpec SparkAngleType() :
1111+
{
1112+
SqlTypeNameSpec t;
1113+
final SqlDataTypeSpec elem;
1114+
final SqlDataTypeSpec key;
1115+
final SqlDataTypeSpec val;
1116+
final SqlTypeName scalar;
1117+
final Span s = Span.of();
1118+
}
1119+
{
1120+
(
1121+
<ARRAY> <LT> elem = SparkDataType() <GT> {
1122+
t = new SqlCollectionTypeNameSpec(
1123+
elem.getTypeNameSpec(), SqlTypeName.ARRAY, getPos());
1124+
}
1125+
|
1126+
<MAP> <LT> key = SparkDataType() <COMMA> val = SparkDataType() <GT> {
1127+
t = new SqlMapTypeNameSpec(key, val, getPos());
1128+
}
1129+
|
1130+
// Spark scalar aliases (string/long/short/byte/bool/timestamp_ltz/timestamp_ntz) in ANY
1131+
// DataType position -- e.g. CAST(x AS STRING) -- not only on the '::' RHS (where
1132+
// SparkTypeName already accepts them). SparkScalarAlias maps STRING->VARCHAR, LONG->BIGINT,
1133+
// etc. (line ~1041). This is the native, nested-paren-safe replacement for the regex
1134+
// CAST(.. AS STRING)->CAST(.. AS VARCHAR) rewrite that lived in SparkSqlPreprocessor.
1135+
// Guarded by an inlined boolean LOOKAHEAD (mirroring SparkTypeName's branch) so it runs
1136+
// inside the generated jj_3R_* lookahead methods, which are not declared to throw
1137+
// ParseException; the wrapping core TypeName() LOOKAHEAD(2) only commits here when token 1
1138+
// is one of these aliases (an IDENTIFIER, never a keyword core type).
1139+
LOOKAHEAD({ getToken(1).kind == IDENTIFIER && getToken(1).image != null
1140+
&& (getToken(1).image.equalsIgnoreCase("STRING")
1141+
|| getToken(1).image.equalsIgnoreCase("LONG")
1142+
|| getToken(1).image.equalsIgnoreCase("SHORT")
1143+
|| getToken(1).image.equalsIgnoreCase("BYTE")
1144+
|| getToken(1).image.equalsIgnoreCase("BOOL")
1145+
|| getToken(1).image.equalsIgnoreCase("TIMESTAMP_LTZ")
1146+
|| getToken(1).image.equalsIgnoreCase("TIMESTAMP_NTZ")) })
1147+
scalar = SparkScalarAlias(s) {
1148+
t = new SqlBasicTypeNameSpec(scalar, s.end(this));
1149+
}
1150+
|
1151+
t = SparkStructType()
1152+
)
1153+
{
1154+
return t;
1155+
}
1156+
}
1157+
9021158
// End parserImpls.ftl

‎sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/planner/BeamRuleSets.java‎

Lines changed: 14 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,12 @@ public class BeamRuleSets {
6060
ImmutableList.of(
6161
// Rules for window functions
6262
CoreRules.PROJECT_TO_LOGICAL_PROJECT_AND_WINDOW,
63+
// Extract an OVER nested inside a Calc expression (e.g. CASE(isnull(x), first_value(x)
64+
// IGNORE NULLS OVER (...), x) as emitted by pandas ffill / bfill) into a LogicalWindow.
65+
// Without this, an OVER that lands inside a Calc -- after PROJECT_TO_CALC converts the
66+
// enclosing project -- stays buried in a BeamCalcRel, which cannot evaluate window
67+
// functions, leaving the plan with no finite-cost alternative.
68+
CoreRules.CALC_TO_WINDOW,
6369
// Rules so we only have to implement Calc
6470
CoreRules.FILTER_TO_CALC,
6571
CoreRules.PROJECT_TO_CALC,
@@ -151,13 +157,14 @@ public class BeamRuleSets {
151157
ImmutableList.of(BeamEnumerableConverterRule.INSTANCE);
152158

153159
public static Collection<RuleSet> getRuleSets() {
160+
return ImmutableList.of(RuleSets.ofList(getAllRules()));
161+
}
154162

155-
return ImmutableList.of(
156-
RuleSets.ofList(
157-
ImmutableList.<RelOptRule>builder()
158-
.addAll(BEAM_CONVERTERS)
159-
.addAll(BEAM_TO_ENUMERABLE)
160-
.addAll(LOGICAL_OPTIMIZATIONS)
161-
.build()));
163+
public static List<RelOptRule> getAllRules() {
164+
return ImmutableList.<RelOptRule>builder()
165+
.addAll(BEAM_CONVERTERS)
166+
.addAll(BEAM_TO_ENUMERABLE)
167+
.addAll(LOGICAL_OPTIMIZATIONS)
168+
.build();
162169
}
163170
}

0 commit comments

Comments
 (0)