Implement Spark SQL parser extensions and enhance analytic functions

Author: damccorm

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<f:T,..>) so
+        # they parse anywhere Calcite accepts a DataType (e.g. CAST(x AS array<int>)), 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 <scope>" 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

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<Object> list, ExprContext exprContext, Span s) :
+{
+    final SqlDataTypeSpec dt;
+}
+{
+    <DOUBLE_COLON> {
+        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<array<int>> closes via two GT tokens (no ">>" shift token exists).
+        <ARRAY> <LT> elem = SparkDataType() <GT> {
+            t = new SqlCollectionTypeNameSpec(
+                    elem.getTypeNameSpec(), SqlTypeName.ARRAY, getPos());
+        }
+    |
+        // Spark MAP<K,V> with Spark-aware inner types (core MapTypeName recurses into core
+        // DataType, which rejects e.g. map<string,int>). LA(2) so a bare MAP without '<' falls
+        // through to the core handler.
+        LOOKAHEAD(2)
+        <MAP> <LT> key = SparkDataType() <COMMA> val = SparkDataType() <GT> {
+            t = new SqlMapTypeNameSpec(key, val, getPos());
+        }
+    |
+        // Spark STRUCT<f:T, ...>. 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<SqlIdentifier> names = new ArrayList<SqlIdentifier>();
+    final List<SqlDataTypeSpec> types = new ArrayList<SqlDataTypeSpec>();
+    SqlIdentifier n;
+    SqlDataTypeSpec ft;
+}
+{
+    // "STRUCT" (guaranteed IDENTIFIER by the LOOKAHEAD at the call site).
+    <IDENTIFIER>
+    <LT>
+        n = SimpleIdentifier() <COLON> ft = SparkDataType() {
+            names.add(n);
+            types.add(ft);
+        }
+        (
+            <COMMA> n = SimpleIdentifier() <COLON> ft = SparkDataType() {
+                names.add(n);
+                types.add(ft);
+            }
+        )*
+    <GT>
+    {
+        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.
+    <IDENTIFIER> {
+        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).
+ *
+ * <p>FLOAT and DOUBLE are RESERVED keyword tokens, so the lexer emits <FLOAT>/<DOUBLE> (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))).
+ *
+ * <p>Hooked in via the `builtinFunctionCallMethods` FreeMarker list, so this is one alternative of
+ * the generated BuiltinFunctionCall(); <FLOAT>/<DOUBLE> are unique first tokens there, so there is
+ * no choice conflict.
+ */
+SqlNode SparkTypeConstructorCall() :
+{
+    final Span s;
+    final SqlNode e;
+    final SqlTypeName tn;
+}
+{
+    ( <FLOAT> { tn = SqlTypeName.FLOAT; } | <DOUBLE> { tn = SqlTypeName.DOUBLE; } )
+    { s = span(); }
+    <LPAREN> e = Expression(ExprContext.ACCEPT_SUB_QUERY) <RPAREN>
+    {
+        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<E>, map<K,V>,
+ * struct<f:T,..>. 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<int>)
+ * 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<string>, map<string,int>).
+ *
+ * <p>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();
+}
+{
+    (
+        <ARRAY> <LT> elem = SparkDataType() <GT> {
+            t = new SqlCollectionTypeNameSpec(
+                    elem.getTypeNameSpec(), SqlTypeName.ARRAY, getPos());
+        }
+    |
+        <MAP> <LT> key = SparkDataType() <COMMA> val = SparkDataType() <GT> {
+            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

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<RuleS
         .ruleSets(ruleSets.toArray(new RuleSet[0]))
         .costFactory(BeamCostModel.FACTORY)
         .typeSystem(connection.getTypeFactory().getTypeSystem())
-        .operatorTable(SqlOperatorTables.chain(opTab0, catalogReader))
+        .operatorTable(
+            SqlOperatorTables.chain(
+                opTab0,
+                // The Spark/PostgreSQL infix cast `expr::TYPE` parses (via the InfixCast grammar
+                // production) to a SqlCall on SqlLibraryOperators.INFIX_CAST, so the validator's
+                // operator table must contain it (the standard table does not). The Spark call-form
+                // collection constructors `ARRAY(...)` / `MAP(...)` likewise parse (via the core
+                // ArrayConstructor/MapConstructor productions) to SqlLibraryOperators.ARRAY/.MAP,
+                // which are SPARK-library operators absent from the standard table; register them
+                // so validation resolves the call form (the bracket form `ARRAY[...]` uses the
+                // standard value-constructors and needs no registration).
+                SqlOperatorTables.of(
+                    org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.fun
+                        .SqlLibraryOperators.INFIX_CAST,
+                    org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.fun
+                        .SqlLibraryOperators.ARRAY,
+                    org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.fun
+                        .SqlLibraryOperators.MAP),
+                catalogReader))
         .sqlToRelConverterConfig(sqlToRelConfig)
         .build();
   }

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

@@ -60,6 +60,12 @@ public class BeamRuleSets {
       ImmutableList.of(
           // Rules for window functions
           CoreRules.PROJECT_TO_LOGICAL_PROJECT_AND_WINDOW,
+          // Extract an OVER nested inside a Calc expression (e.g. CASE(isnull(x), first_value(x)
+          // IGNORE NULLS OVER (...), x) as emitted by pandas ffill / bfill) into a LogicalWindow.
+          // Without this, an OVER that lands inside a Calc -- after PROJECT_TO_CALC converts the
+          // enclosing project -- stays buried in a BeamCalcRel, which cannot evaluate window
+          // functions, leaving the plan with no finite-cost alternative.
+          CoreRules.CALC_TO_WINDOW,
           // Rules so we only have to implement Calc
           CoreRules.FILTER_TO_CALC,
           CoreRules.PROJECT_TO_CALC,
@@ -151,13 +157,14 @@ public class BeamRuleSets {
       ImmutableList.of(BeamEnumerableConverterRule.INSTANCE);
 
   public static Collection<RuleSet> getRuleSets() {
+    return ImmutableList.of(RuleSets.ofList(getAllRules()));
+  }
 
-    return ImmutableList.of(
-        RuleSets.ofList(
-            ImmutableList.<RelOptRule>builder()
-                .addAll(BEAM_CONVERTERS)
-                .addAll(BEAM_TO_ENUMERABLE)
-                .addAll(LOGICAL_OPTIMIZATIONS)
-                .build()));
+  public static List<RelOptRule> getAllRules() {
+    return ImmutableList.<RelOptRule>builder()
+        .addAll(BEAM_CONVERTERS)
+        .addAll(BEAM_TO_ENUMERABLE)
+        .addAll(LOGICAL_OPTIMIZATIONS)
+        .build();
   }
 }