diff --git 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 index d84783118bbd..6da6d6279b3b 100644 --- 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 @@ -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 @@ public String explain(String sqlString) throws ParseException { } } + /** + * 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 = diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/BeamSqlEnv.java.rej b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/BeamSqlEnv.java.rej new file mode 100644 index 000000000000..08dbe6af0d5a --- /dev/null +++ b/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(); 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 aa6f4d121871..ac1140af895e 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 @@ -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 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 props = sqlOptions.getCalciteConnectionProperties(); + if (props == null) { + return config.conformance(); + } + String value = null; + for (Map.Entry 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. diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/CalciteQueryPlanner.java.rej b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/CalciteQueryPlanner.java.rej new file mode 100644 index 000000000000..abe65beb9b8d --- /dev/null +++ b/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; diff --git 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 index 0f0f8970a3ab..9d4bc7930e43 100644 --- 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 @@ -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; diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/QueryPlanner.java.rej b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/QueryPlanner.java.rej new file mode 100644 index 000000000000..bae9825b9cc3 --- /dev/null +++ b/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; + diff --git a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/BeamSqlEnvTest.java b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/BeamSqlEnvTest.java index 3b6dda010427..feafcb4f3660 100644 --- a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/BeamSqlEnvTest.java +++ b/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()); + } }