Add parser conformance resolution and logical plan parsing

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

@@ -47,6 +47,7 @@
 import org.apache.beam.sdk.transforms.SerializableFunction;
 import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.jdbc.CalcitePrepare;
 import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.RelOptUtil;
+import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.RelNode;
 import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.schema.Function;
 import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlKind;
 import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.tools.RuleSet;
@@ -140,6 +141,24 @@
     }
   }
 
+  /**
+   * Parses and validates {@code query} and returns the logical {@link RelNode} (Calcite convention)
+   * without Beam conversion, so callers can rewrite the plan before {@link
+   * #convertToBeamRel(RelNode, QueryParameters)}.
+   */
+  public RelNode parseLogicalPlan(String query) throws ParseException {
+    return planner.parseToRel(query, QueryParameters.ofNone());
+  }
+
+  /**
+   * Registers a Calcite {@link Function} under {@code name} in the live (current) session schema,
+   * so that subsequent SQL parsing/validation resolves {@code name(...)}. Mirrors how user-defined
+   * functions are added at build time.
+   */
+  public void registerSchemaFunction(String name, Function function) {
+    connection.getCurrentSchemaPlus().add(name, function);
+  }
+
   /** BeamSqlEnv's Builder. */
   public static class BeamSqlEnvBuilder {
     private static final String CALCITE_PLANNER =

sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/BeamSqlEnv.java.rej

@@ -0,0 +1,26 @@
+diff a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/BeamSqlEnv.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/BeamSqlEnv.java	(rejected hunks)
+@@ -135,6 +135,24 @@ public class BeamSqlEnv {
+     return planner.convertToBeamRel(relNode, QueryParameters.ofNone());
+   }
+
++  /**
++   * Parses and validates {@code query} and returns the logical {@link RelNode} (Calcite convention)
++   * without Beam conversion, so callers can rewrite the plan before {@link
++   * #convertToBeamRel(RelNode)}.
++   */
++  public RelNode parseLogicalPlan(String query) throws ParseException {
++    return planner.parseToRel(query, QueryParameters.ofNone());
++  }
++
++  /**
++   * Registers a Calcite {@link Function} under {@code name} in the live (current) session schema,
++   * so that subsequent SQL parsing/validation resolves {@code name(...)}. Mirrors how user-defined
++   * functions are added at build time.
++   */
++  public void registerSchemaFunction(String name, Function function) {
++    connection.getCurrentSchemaPlus().add(name, function);
++  }
++
+   public org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlOperatorTable
+       getOperatorTable() {
+     return planner.getOperatorTable();

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

@@ -31,6 +31,7 @@
 import org.apache.beam.sdk.extensions.sql.impl.rel.BeamRelNode;
 import org.apache.beam.sdk.extensions.sql.impl.rel.BeamSqlRelUtils;
 import org.apache.beam.sdk.extensions.sql.impl.udf.BeamBuiltinFunctionProvider;
+import org.apache.beam.sdk.options.PipelineOptions;
 import org.apache.beam.vendor.calcite.v1_40_0.com.google.common.collect.Table;
 import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.config.CalciteConnectionConfig;
 import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.jdbc.CalciteSchema;
@@ -64,6 +65,8 @@
 import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.parser.SqlParser;
 import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.parser.SqlParserImplFactory;
 import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.util.SqlOperatorTables;
+import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.validate.SqlConformance;
+import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.validate.SqlConformanceEnum;
 import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql2rel.SqlToRelConverter;
 import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.tools.FrameworkConfig;
 import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.tools.Frameworks;
@@ -120,12 +123,20 @@ private void loadBuiltinFunctions(JdbcConnection jdbcConnection) {
 
   public FrameworkConfig defaultConfig(JdbcConnection connection, Collection<RuleSet> ruleSets) {
     final CalciteConnectionConfig config = connection.config();
+    // Resolve the parser conformance. Calcite's Avatica JDBC connect path silently drops the
+    // {@code conformance} connection property (it is not in the driver's registered property set),
+    // so {@code config.conformance()} is always DEFAULT here even when callers set it via
+    // {@code BeamSqlPipelineOptions.calciteConnectionProperties}. We therefore read that map
+    // directly from the connection's pipeline options and let it override. This keeps the behavior
+    // opt-in: with no {@code conformance} property the connection's own (DEFAULT) value is used, so
+    // existing Beam SQL behavior is unchanged.
+    final SqlConformance conformance = resolveConformance(connection, config);
     final SqlParser.ConfigBuilder parserConfig =
         SqlParser.configBuilder()
             .setQuotedCasing(config.quotedCasing())
             .setUnquotedCasing(config.unquotedCasing())
             .setQuoting(config.quoting())
-            .setConformance(config.conformance())
+            .setConformance(conformance)
             .setCaseSensitive(config.caseSensitive());
     final SqlParserImplFactory parserFactory =
         config.parserFactory(SqlParserImplFactory.class, null);
@@ -163,6 +174,43 @@ public FrameworkConfig defaultConfig(JdbcConnection connection, Collection<RuleS
         .build();
   }
 
+  /**
+   * Resolves the {@link SqlConformance} for the parser. Prefers an explicit {@code conformance}
+   * entry in {@link BeamSqlPipelineOptions#getCalciteConnectionProperties()} (looked up
+   * case-insensitively, value matched against {@link SqlConformanceEnum}); otherwise falls back to
+   * the connection's own conformance. This is the bridge for {@code conformance=BABEL}, which the
+   * Avatica JDBC connect path drops, enabling Spark-SQL spellings (e.g. the {@code !=} operator)
+   * the default conformance rejects. Returns the connection default on any unrecognized value.
+   */
+  private static SqlConformance resolveConformance(
+      JdbcConnection connection, CalciteConnectionConfig config) {
+    PipelineOptions options = connection.getPipelineOptions();
+    if (options == null) {
+      return config.conformance();
+    }
+    BeamSqlPipelineOptions sqlOptions = options.as(BeamSqlPipelineOptions.class);
+    Map<String, String> props = sqlOptions.getCalciteConnectionProperties();
+    if (props == null) {
+      return config.conformance();
+    }
+    String value = null;
+    for (Map.Entry<String, String> e : props.entrySet()) {
+      if ("conformance".equalsIgnoreCase(e.getKey())) {
+        value = e.getValue();
+        break;
+      }
+    }
+    if (value == null) {
+      return config.conformance();
+    }
+    try {
+      return SqlConformanceEnum.valueOf(value.trim().toUpperCase());
+    } catch (IllegalArgumentException ex) {
+      LOG.warn("Unrecognized calcite conformance '{}', using {}", value, config.conformance());
+      return config.conformance();
+    }
+  }
+
   /** Parse input SQL query, and return a {@link SqlNode} as grammar tree. */
   @Override
   public SqlNode parse(String sqlStatement) throws ParseException {
@@ -177,6 +225,29 @@ public SqlNode parse(String sqlStatement) throws ParseException {
     return parsed;
   }
 
+  @Override
+  public RelNode parseToRel(String sqlStatement, QueryParameters queryParameters)
+      throws ParseException, SqlConversionException {
+    Preconditions.checkArgument(
+        queryParameters.getKind() == Kind.NONE,
+        "Beam SQL Calcite dialect does not yet support query parameters.");
+    try {
+      SqlNode parsed = planner.parse(sqlStatement);
+      TableResolutionUtils.setupCustomTableResolution(connection, parsed);
+      SqlNode validated = planner.validate(parsed);
+      // root of original logical plan
+      RelRoot root = planner.rel(validated);
+      return root.rel;
+    } catch (RelConversionException e) {
+      throw new SqlConversionException(
+          String.format("Unable to convert query %s", sqlStatement), e);
+    } catch (SqlParseException | ValidationException e) {
+      throw new ParseException(String.format("Unable to parse query %s", sqlStatement), e);
+    } finally {
+      planner.close();
+    }
+  }
+
   /**
    * It parses and validate the input query, then convert into a {@link BeamRelNode} tree. Note that
    * query parameters are not yet supported.

sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/CalciteQueryPlanner.java.rej

@@ -0,0 +1,31 @@
+diff 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	(rejected hunks)
+@@ -268,6 +316,29 @@ public class CalciteQueryPlanner implements QueryPlanner {
+     return newRel.copy(newRel.getTraitSet(), newInputs);
+   }
+
++  @Override
++  public RelNode parseToRel(String sqlStatement, QueryParameters queryParameters)
++      throws ParseException, SqlConversionException {
++    Preconditions.checkArgument(
++        queryParameters.getKind() == Kind.NONE,
++        "Beam SQL Calcite dialect does not yet support query parameters.");
++    try {
++      SqlNode parsed = planner.parse(sqlStatement);
++      TableResolutionUtils.setupCustomTableResolution(connection, parsed);
++      SqlNode validated = planner.validate(parsed);
++      // root of original logical plan
++      RelRoot root = planner.rel(validated);
++      return root.rel;
++    } catch (RelConversionException e) {
++      throw new SqlConversionException(
++          String.format("Unable to convert query %s", sqlStatement), e);
++    } catch (SqlParseException | ValidationException e) {
++      throw new ParseException(String.format("Unable to parse query %s", sqlStatement), e);
++    } finally {
++      planner.close();
++    }
++  }
++
+   @Override
+   public BeamRelNode convertToBeamRel(RelNode relNode, QueryParameters queryParameters) {
+     RelNode beamRelNode;

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

@@ -22,6 +22,7 @@
 import java.util.List;
 import java.util.Map;
 import org.apache.beam.sdk.extensions.sql.impl.rel.BeamRelNode;
+import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.RelNode;
 import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlNode;
 import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.tools.RuleSet;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
@@ -39,6 +40,18 @@ public interface QueryPlanner {
   BeamRelNode convertToBeamRel(String sqlStatement, QueryParameters queryParameters)
       throws ParseException, SqlConversionException;
 
+  /**
+   * Parses and validates {@code sqlStatement} and returns the logical {@link RelNode} (standard
+   * Calcite convention), WITHOUT converting to Beam physical rels. This lets callers rewrite the
+   * logical plan (e.g. inject custom rels) before {@link #convertToBeamRel(String,
+   * QueryParameters)}.
+   */
+  default RelNode parseToRel(String sqlStatement, QueryParameters queryParameters)
+      throws ParseException, SqlConversionException {
+    throw new UnsupportedOperationException(
+        "parseToRel is not implemented by " + getClass().getName());
+  }
+
   /** Parse input SQL query, and return a {@link SqlNode} as grammar tree. */
   SqlNode parse(String sqlStatement) throws ParseException;
 

sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/QueryPlanner.java.rej

@@ -0,0 +1,20 @@
+diff a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/QueryPlanner.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/QueryPlanner.java	(rejected hunks)
+@@ -45,6 +45,18 @@ public interface QueryPlanner {
+   BeamRelNode convertToBeamRel(RelNode sqlStatement, QueryParameters queryParameters)
+       throws SqlConversionException;
+
++  /**
++   * Parses and validates {@code sqlStatement} and returns the logical {@link RelNode} (standard
++   * Calcite convention), WITHOUT converting to Beam physical rels. This lets callers rewrite the
++   * logical plan (e.g. inject custom rels) before {@link #convertToBeamRel(RelNode,
++   * QueryParameters)}.
++   */
++  default RelNode parseToRel(String sqlStatement, QueryParameters queryParameters)
++      throws ParseException, SqlConversionException {
++    throw new UnsupportedOperationException(
++        "parseToRel is not implemented by " + getClass().getName());
++  }
++
+   /** Parse input SQL query, and return a {@link SqlNode} as grammar tree. */
+   SqlNode parse(String sqlStatement) throws ParseException;
+

sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/BeamSqlEnvTest.java

@@ -67,4 +67,21 @@ public void testPlannerClassNotFound() {
         .setPipelineOptions(PipelineOptionsFactory.create())
         .build();
   }
+
+  @Test
+  public void testParseLogicalPlan() throws Exception {
+    TestTableProvider root = new TestTableProvider();
+    BeamSqlEnv env =
+        BeamSqlEnv.builder(root).setPipelineOptions(PipelineOptionsFactory.create()).build();
+    Connection connection = env.connection;
+    connection.createStatement().execute("CREATE EXTERNAL TABLE person (id INT) TYPE test");
+
+    org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.RelNode logicalPlan =
+        env.parseLogicalPlan("SELECT id FROM person");
+
+    // the logical plan should be a Project
+    assertEquals(
+        "org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.logical.LogicalProject",
+        logicalPlan.getClass().getName());
+  }
 }