diff --git
a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/SqlTransform.java
b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/SqlTransform.java
index 5fd08806c525..f5613c9b56bf 100644
---
a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/SqlTransform.java
+++
b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/SqlTransform.java
@@ -19,12 +19,15 @@
import com.google.auto.value.AutoValue;
import java.util.Collections;
+import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import javax.annotation.Nullable;
import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.extensions.sql.impl.BeamSqlEnv;
import org.apache.beam.sdk.extensions.sql.impl.rel.BeamSqlRelUtils;
import org.apache.beam.sdk.extensions.sql.impl.schema.BeamPCollectionTable;
+import org.apache.beam.sdk.extensions.sql.meta.provider.TableProvider;
import org.apache.beam.sdk.transforms.Combine;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.SerializableFunction;
@@ -86,9 +89,17 @@
abstract boolean autoUdfUdafLoad();
+ abstract Map<String, TableProvider> tableProviderMap();
+
+ abstract @Nullable String defaultTableProvider();
+
@Override
public PCollection<Row> expand(PInput input) {
BeamSqlEnv sqlEnv = BeamSqlEnv.readOnly(PCOLLECTION_NAME,
toTableMap(input));
+ tableProviderMap().forEach(sqlEnv::addSchema);
+ if (defaultTableProvider() != null) {
+ sqlEnv.setCurrentSchema(defaultTableProvider());
+ }
// TODO: validate duplicate functions.
sqlEnv.loadBeamBuiltinFunctions();
@@ -154,10 +165,21 @@ public static SqlTransform query(String queryString) {
.setQueryString(queryString)
.setUdafDefinitions(Collections.emptyList())
.setUdfDefinitions(Collections.emptyList())
+ .setTableProviderMap(Collections.emptyMap())
.setAutoUdfUdafLoad(false)
.build();
}
+ public SqlTransform withTableProvider(String name, TableProvider
tableProvider) {
+ Map<String, TableProvider> map = new HashMap<>(tableProviderMap());
+ map.put(name, tableProvider);
+ return toBuilder().setTableProviderMap(ImmutableMap.copyOf(map)).build();
+ }
+
+ public SqlTransform withDefaultTableProvider(String name, TableProvider
tableProvider) {
+ return withTableProvider(name,
tableProvider).toBuilder().setDefaultTableProvider(name).build();
+ }
+
public SqlTransform withAutoUdfUdafLoad(boolean autoUdfUdafLoad) {
return toBuilder().setAutoUdfUdafLoad(autoUdfUdafLoad).build();
}
@@ -215,6 +237,10 @@ static Builder builder() {
abstract Builder setAutoUdfUdafLoad(boolean autoUdfUdafLoad);
+ abstract Builder setTableProviderMap(Map<String, TableProvider>
tableProviderMap);
+
+ abstract Builder setDefaultTableProvider(@Nullable String
defaultTableProvider);
+
abstract SqlTransform build();
}
diff --git
a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/BeamQueryPlanner.java
b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/BeamQueryPlanner.java
index 65b2f1556555..83993600f941 100644
---
a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/BeamQueryPlanner.java
+++
b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/BeamQueryPlanner.java
@@ -54,9 +54,13 @@
class BeamQueryPlanner {
private static final Logger LOG =
LoggerFactory.getLogger(BeamQueryPlanner.class);
- private final FrameworkConfig config;
+ private JdbcConnection connection;
BeamQueryPlanner(JdbcConnection connection) {
+ this.connection = connection;
+ }
+
+ public FrameworkConfig config() {
final CalciteConnectionConfig config = connection.config();
final SqlParser.ConfigBuilder parserConfig =
SqlParser.configBuilder()
@@ -85,17 +89,16 @@
final SqlOperatorTable opTab0 =
connection.config().fun(SqlOperatorTable.class,
SqlStdOperatorTable.instance());
- this.config =
- Frameworks.newConfigBuilder()
- .parserConfig(parserConfig.build())
- .defaultSchema(defaultSchema)
- .traitDefs(traitDefs)
- .context(Contexts.of(connection.config()))
- .ruleSets(BeamRuleSets.getRuleSets())
- .costFactory(null)
- .typeSystem(connection.getTypeFactory().getTypeSystem())
- .operatorTable(ChainedSqlOperatorTable.of(opTab0, catalogReader))
- .build();
+ return Frameworks.newConfigBuilder()
+ .parserConfig(parserConfig.build())
+ .defaultSchema(defaultSchema)
+ .traitDefs(traitDefs)
+ .context(Contexts.of(connection.config()))
+ .ruleSets(BeamRuleSets.getRuleSets())
+ .costFactory(null)
+ .typeSystem(connection.getTypeFactory().getTypeSystem())
+ .operatorTable(ChainedSqlOperatorTable.of(opTab0, catalogReader))
+ .build();
}
/** Parse input SQL query, and return a {@link SqlNode} as grammar tree. */
@@ -141,6 +144,6 @@ public BeamRelNode convertToBeamRel(String sqlStatement)
}
private Planner getPlanner() {
- return Frameworks.getPlanner(config);
+ return Frameworks.getPlanner(config());
}
}
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 5841efb64f8c..e1589bdea81e 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
@@ -18,6 +18,7 @@
package org.apache.beam.sdk.extensions.sql.impl;
import java.lang.reflect.Method;
+import java.sql.SQLException;
import java.util.List;
import java.util.Map;
import java.util.ServiceLoader;
@@ -80,6 +81,18 @@ private void registerBuiltinUdf(Map<String, List<Method>>
methods) {
}
}
+ public void addSchema(String name, TableProvider tableProvider) {
+ connection.setSchema(name, tableProvider);
+ }
+
+ public void setCurrentSchema(String name) {
+ try {
+ connection.setSchema(name);
+ } catch (SQLException e) {
+ throw new RuntimeException(e);
+ }
+ }
+
/** Register a UDF function which can be used in SQL expression. */
public void registerUdf(String functionName, Class<?> clazz, String method) {
connection.getCurrentSchemaPlus().add(functionName, UdfImpl.create(clazz,
method));
diff --git
a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/JdbcConnection.java
b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/JdbcConnection.java
index a4a7f479c10b..f3327d857347 100644
---
a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/JdbcConnection.java
+++
b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/JdbcConnection.java
@@ -64,6 +64,7 @@ private JdbcConnection(CalciteConnection connection) throws
SQLException {
JdbcConnection jdbcConnection = new JdbcConnection(connection);
jdbcConnection.setPipelineOptionsMap(extractPipelineOptions(connection));
+ jdbcConnection.getRootSchema().setCacheEnabled(false);
jdbcConnection.setSchema(
connection.getSchema(),
BeamCalciteSchemaFactory.fromInitialEmptySchema(jdbcConnection));
return jdbcConnection;
diff --git
a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlMultipleSchemasTest.java
b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlMultipleSchemasTest.java
new file mode 100644
index 000000000000..5d9e5bf2254c
--- /dev/null
+++
b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlMultipleSchemasTest.java
@@ -0,0 +1,253 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.extensions.sql;
+
+import com.google.common.collect.ImmutableMap;
+import java.util.Arrays;
+import org.apache.beam.sdk.extensions.sql.impl.schema.BeamPCollectionTable;
+import org.apache.beam.sdk.extensions.sql.meta.provider.ReadOnlyTableProvider;
+import org.apache.beam.sdk.extensions.sql.meta.provider.TableProvider;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.testing.PAssert;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.values.PBegin;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.ExpectedException;
+
+/** Ad-hoc tests for CAST. */
+public class BeamSqlMultipleSchemasTest {
+
+ private static final Schema ROW_SCHEMA =
+
Schema.builder().addInt32Field("f_int").addStringField("f_string").build();
+
+ @Rule public final TestPipeline pipeline = TestPipeline.create();
+ @Rule public ExpectedException exceptions = ExpectedException.none();
+
+ @Test
+ public void testSelectFromQualifiedPCollection() {
+ PCollection<Row> input = pipeline.apply(create(row(1, "strstr")));
+
+ PCollection<Row> result =
+ input.apply(SqlTransform.query("SELECT f_int, f_string FROM
beam.PCOLLECTION"));
+
+ PAssert.that(result).containsInAnyOrder(row(1, "strstr"));
+ pipeline.run();
+ }
+
+ @Test
+ public void testSelectFromUnqualifiedPCollection() {
+ PCollection<Row> input = pipeline.apply(create(row(1, "strstr")));
+
+ PCollection<Row> result =
+ input.apply(SqlTransform.query("SELECT f_int, f_string FROM
PCOLLECTION"));
+
+ PAssert.that(result).containsInAnyOrder(row(1, "strstr"));
+ pipeline.run();
+ }
+
+ @Test
+ public void testSelectFromExtraSchema() {
+ PCollection<Row> inputMain =
+ pipeline.apply("mainInput", create(row(1, "pcollection_1"), row(2,
"pcollection_2")));
+
+ PCollection<Row> inputExtra =
+ pipeline.apply("extraInput", create(row(1, "_extra_table_1"), row(2,
"_extra_table_2")));
+
+ TableProvider extraInputProvider = extraTableProvider("extraTable",
inputExtra);
+
+ PCollection<Row> result =
+ inputMain.apply(
+ SqlTransform.query("SELECT f_int, f_string FROM
extraSchema.extraTable")
+ .withTableProvider("extraSchema", extraInputProvider));
+
+ PAssert.that(result).containsInAnyOrder(row(1, "_extra_table_1"), row(2,
"_extra_table_2"));
+ pipeline.run();
+ }
+
+ @Test
+ public void testOverrideUnqualifiedMainSchema() {
+ PCollection<Row> inputMain =
+ pipeline.apply("mainInput", create(row(1, "pcollection_1"), row(2,
"pcollection_2")));
+
+ PCollection<Row> inputExtra =
+ pipeline.apply("extraInput", create(row(1, "_extra_table_1"), row(2,
"_extra_table_2")));
+
+ TableProvider extraInputProvider = extraTableProvider("extraTable",
inputExtra);
+
+ PCollection<Row> result =
+ inputMain.apply(
+ SqlTransform.query("SELECT f_int, f_string FROM extraTable")
+ .withTableProvider("beam", extraInputProvider));
+
+ PAssert.that(result).containsInAnyOrder(row(1, "_extra_table_1"), row(2,
"_extra_table_2"));
+ pipeline.run();
+ }
+
+ @Test
+ public void testOverrideQualifiedMainSchema() {
+ PCollection<Row> inputMain =
+ pipeline.apply("mainInput", create(row(1, "pcollection_1"), row(2,
"pcollection_2")));
+
+ PCollection<Row> inputExtra =
+ pipeline.apply("extraInput", create(row(1, "_extra_table_1"), row(2,
"_extra_table_2")));
+
+ TableProvider extraInputProvider = extraTableProvider("extraTable",
inputExtra);
+
+ PCollection<Row> result =
+ inputMain.apply(
+ SqlTransform.query("SELECT f_int, f_string FROM beam.extraTable")
+ .withTableProvider("beam", extraInputProvider));
+
+ PAssert.that(result).containsInAnyOrder(row(1, "_extra_table_1"), row(2,
"_extra_table_2"));
+ pipeline.run();
+ }
+
+ @Test
+ public void testSetDefaultUnqualifiedSchema() {
+ PCollection<Row> inputMain =
+ pipeline.apply("mainInput", create(row(1, "pcollection_1"), row(2,
"pcollection_2")));
+
+ PCollection<Row> inputExtra =
+ pipeline.apply("extraInput", create(row(1, "_extra_table_1"), row(2,
"_extra_table_2")));
+
+ TableProvider extraInputProvider = extraTableProvider("extraTable",
inputExtra);
+
+ PCollection<Row> result =
+ inputMain.apply(
+ SqlTransform.query("SELECT f_int, f_string FROM extraTable")
+ .withDefaultTableProvider("extraSchema", extraInputProvider));
+
+ PAssert.that(result).containsInAnyOrder(row(1, "_extra_table_1"), row(2,
"_extra_table_2"));
+ pipeline.run();
+ }
+
+ @Test
+ public void testSetDefaultUnqualifiedSchemaAndJoin() {
+ PCollection<Row> inputMain =
+ pipeline.apply("mainInput", create(row(1, "pcollection_1"), row(2,
"pcollection_2")));
+
+ PCollection<Row> inputExtra =
+ pipeline.apply("extraInput", create(row(1, "_extra_table_1"), row(2,
"_extra_table_2")));
+
+ TableProvider extraInputProvider = extraTableProvider("extraTable",
inputExtra);
+
+ PCollection<Row> result =
+ inputMain.apply(
+ SqlTransform.query(
+ "SELECT extra.f_int, main.f_string || extra.f_string AS
f_string \n"
+ + "FROM extraTable AS extra \n"
+ + " INNER JOIN \n"
+ + " beam.PCOLLECTION AS main \n"
+ + " ON main.f_int = extra.f_int")
+ .withDefaultTableProvider("extraSchema", extraInputProvider));
+
+ PAssert.that(result)
+ .containsInAnyOrder(
+ row(1, "pcollection_1_extra_table_1"), row(2,
"pcollection_2_extra_table_2"));
+ pipeline.run();
+ }
+
+ @Test
+ public void testSetDefaultQualifiedSchema() {
+ PCollection<Row> inputMain =
+ pipeline.apply("mainInput", create(row(1, "pcollection_1"), row(2,
"pcollection_2")));
+
+ PCollection<Row> inputExtra =
+ pipeline.apply("extraInput", create(row(1, "_extra_table_1"), row(2,
"_extra_table_2")));
+
+ TableProvider extraInputProvider = extraTableProvider("extraTable",
inputExtra);
+
+ PCollection<Row> result =
+ inputMain.apply(
+ SqlTransform.query("SELECT f_int, f_string FROM
extraSchema.extraTable")
+ .withDefaultTableProvider("extraSchema", extraInputProvider));
+
+ PAssert.that(result).containsInAnyOrder(row(1, "_extra_table_1"), row(2,
"_extra_table_2"));
+ pipeline.run();
+ }
+
+ @Test
+ public void testJoinWithExtraSchema() {
+ PCollection<Row> inputMain =
+ pipeline.apply("mainInput", create(row(1, "pcollection_1"), row(2,
"pcollection_2")));
+
+ PCollection<Row> inputExtra =
+ pipeline.apply("extraInput", create(row(1, "_extra_table_1"), row(2,
"_extra_table_2")));
+
+ TableProvider extraInputProvider = extraTableProvider("extraTable",
inputExtra);
+
+ PCollection<Row> result =
+ inputMain.apply(
+ SqlTransform.query(
+ "SELECT extra.f_int, main.f_string || extra.f_string AS
f_string \n"
+ + "FROM extraSchema.extraTable AS extra \n"
+ + " INNER JOIN \n"
+ + " PCOLLECTION AS main \n"
+ + " ON main.f_int = extra.f_int")
+ .withTableProvider("extraSchema", extraInputProvider));
+
+ PAssert.that(result)
+ .containsInAnyOrder(
+ row(1, "pcollection_1_extra_table_1"), row(2,
"pcollection_2_extra_table_2"));
+ pipeline.run();
+ }
+
+ @Test
+ public void testJoinQualifiedMainWithExtraSchema() {
+ PCollection<Row> inputMain =
+ pipeline.apply("mainInput", create(row(1, "pcollection_1"), row(2,
"pcollection_2")));
+
+ PCollection<Row> inputExtra =
+ pipeline.apply("extraInput", create(row(1, "_extra_table_1"), row(2,
"_extra_table_2")));
+
+ TableProvider extraInputProvider = extraTableProvider("extraTable",
inputExtra);
+
+ PCollection<Row> result =
+ inputMain.apply(
+ SqlTransform.query(
+ "SELECT extra.f_int, main.f_string || extra.f_string AS
f_string \n"
+ + "FROM extraSchema.extraTable AS extra \n"
+ + " INNER JOIN \n"
+ + " beam.PCOLLECTION AS main \n"
+ + " ON main.f_int = extra.f_int")
+ .withTableProvider("extraSchema", extraInputProvider));
+
+ PAssert.that(result)
+ .containsInAnyOrder(
+ row(1, "pcollection_1_extra_table_1"), row(2,
"pcollection_2_extra_table_2"));
+ pipeline.run();
+ }
+
+ private TableProvider extraTableProvider(String tableName, PCollection<Row>
rows) {
+ return new ReadOnlyTableProvider(
+ "testExtraTableProvider", ImmutableMap.of(tableName, new
BeamPCollectionTable<>(rows)));
+ }
+
+ private Row row(int fIntValue, String fStringValue) {
+ return Row.withSchema(ROW_SCHEMA).addValues(fIntValue,
fStringValue).build();
+ }
+
+ private PTransform<PBegin, PCollection<Row>> create(Row... rows) {
+ return Create.of(Arrays.asList(rows)).withRowSchema(ROW_SCHEMA);
+ }
+}
With regards,
Apache Git Services
