Search in sources :

Example 11 with BeamRelNode

use of org.apache.beam.sdk.extensions.sql.impl.rel.BeamRelNode in project beam by apache.

the class CalciteQueryPlannerTest method testNonCumulativeCostMetadataHandler.

@Test
public void testNonCumulativeCostMetadataHandler() {
    String sql = "select * from medium_table";
    BeamRelNode root = env.parseQuery(sql);
    Assert.assertTrue(root.getCluster().getMetadataQuery().getNonCumulativeCost(root) instanceof BeamCostModel);
    Assert.assertFalse(root.getCluster().getMetadataQuery().getNonCumulativeCost(root).isInfinite());
}
Also used : BeamRelNode(org.apache.beam.sdk.extensions.sql.impl.rel.BeamRelNode) BaseRelTest(org.apache.beam.sdk.extensions.sql.impl.rel.BaseRelTest) Test(org.junit.Test)

Example 12 with BeamRelNode

use of org.apache.beam.sdk.extensions.sql.impl.rel.BeamRelNode in project beam by apache.

the class CalciteQueryPlannerTest method testCumulativeCostMetaDataHandler.

@Test
public void testCumulativeCostMetaDataHandler() {
    // This handler is not our handler. It tests if the cumulative handler of Calcite works as
    // expected.
    String sql = "select * from medium_table";
    BeamRelNode root = env.parseQuery(sql);
    Assert.assertTrue(root.getCluster().getMetadataQuery().getCumulativeCost(root) instanceof BeamCostModel);
    Assert.assertFalse(root.getCluster().getMetadataQuery().getCumulativeCost(root).isInfinite());
}
Also used : BeamRelNode(org.apache.beam.sdk.extensions.sql.impl.rel.BeamRelNode) BaseRelTest(org.apache.beam.sdk.extensions.sql.impl.rel.BaseRelTest) Test(org.junit.Test)

Example 13 with BeamRelNode

use of org.apache.beam.sdk.extensions.sql.impl.rel.BeamRelNode in project beam by apache.

the class IOPushDownRuleTest method testIsProjectRenameOnlyProgram.

@Test
public void testIsProjectRenameOnlyProgram() {
    List<Pair<Pair<String, Boolean>, Boolean>> tests = ImmutableList.of(// Selecting fields in a different order is only allowed with project push-down.
    Pair.of(Pair.of("select unused2, name, id from TEST", true), true), Pair.of(Pair.of("select unused2, name, id from TEST", false), false), Pair.of(Pair.of("select id from TEST", false), true), Pair.of(Pair.of("select * from TEST", false), true), Pair.of(Pair.of("select id, name from TEST", false), true), Pair.of(Pair.of("select id+10 from TEST", false), false), // Note that we only care about projects.
    Pair.of(Pair.of("select id from TEST where name='one'", false), true));
    for (Pair<Pair<String, Boolean>, Boolean> test : tests) {
        String sqlQuery = test.left.left;
        boolean projectPushDownSupported = test.left.right;
        boolean expectedAnswer = test.right;
        BeamRelNode basicRel = sqlEnv.parseQuery(sqlQuery);
        assertThat(basicRel, instanceOf(Calc.class));
        Calc calc = (Calc) basicRel;
        assertThat(test.toString(), BeamIOPushDownRule.INSTANCE.isProjectRenameOnlyProgram(calc.getProgram(), projectPushDownSupported), equalTo(expectedAnswer));
    }
}
Also used : BeamRelNode(org.apache.beam.sdk.extensions.sql.impl.rel.BeamRelNode) Calc(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.core.Calc) Pair(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.util.Pair) Test(org.junit.Test)

Example 14 with BeamRelNode

use of org.apache.beam.sdk.extensions.sql.impl.rel.BeamRelNode in project beam by apache.

the class ZetaSqlJavaUdfTest method testNestedJavaUdf.

@Test
public void testNestedJavaUdf() {
    String sql = String.format("CREATE FUNCTION increment(i INT64) RETURNS INT64 LANGUAGE java " + "OPTIONS (path='%s'); " + "SELECT increment(increment(1));", jarPath);
    ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config);
    BeamRelNode beamRelNode = zetaSQLQueryPlanner.convertToBeamRel(sql);
    PCollection<Row> stream = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
    Schema singleField = Schema.builder().addInt64Field("field1").build();
    PAssert.that(stream).containsInAnyOrder(Row.withSchema(singleField).addValues(3L).build());
    pipeline.run().waitUntilFinish(Duration.standardMinutes(PIPELINE_EXECUTION_WAITTIME_MINUTES));
}
Also used : BeamRelNode(org.apache.beam.sdk.extensions.sql.impl.rel.BeamRelNode) Schema(org.apache.beam.sdk.schemas.Schema) Matchers.containsString(org.hamcrest.Matchers.containsString) Row(org.apache.beam.sdk.values.Row) Test(org.junit.Test)

Example 15 with BeamRelNode

use of org.apache.beam.sdk.extensions.sql.impl.rel.BeamRelNode in project beam by apache.

the class ZetaSqlJavaUdfTest method testUdfFromCatalog.

/**
 * This tests a subset of the code path used by {@link #testSqlTransformRegisterUdf()}.
 */
@Test
public void testUdfFromCatalog() throws NoSuchMethodException {
    // Add IncrementFn to Calcite schema.
    JdbcConnection jdbcConnection = JdbcDriver.connect(new ReadOnlyTableProvider("empty_table_provider", ImmutableMap.of()), PipelineOptionsFactory.create());
    Method method = IncrementFn.class.getMethod("eval", Long.class);
    jdbcConnection.getCurrentSchemaPlus().add("increment", ScalarFunctionImpl.create(method));
    this.config = Frameworks.newConfigBuilder(config).defaultSchema(jdbcConnection.getCurrentSchemaPlus()).build();
    String sql = "SELECT increment(0);";
    ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config);
    BeamRelNode beamRelNode = zetaSQLQueryPlanner.convertToBeamRel(sql);
    PCollection<Row> stream = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
    final Schema schema = Schema.builder().addInt64Field("field1").build();
    PAssert.that(stream).containsInAnyOrder(Row.withSchema(schema).addValues(1L).build());
    pipeline.run().waitUntilFinish(Duration.standardMinutes(PIPELINE_EXECUTION_WAITTIME_MINUTES));
}
Also used : BeamRelNode(org.apache.beam.sdk.extensions.sql.impl.rel.BeamRelNode) Schema(org.apache.beam.sdk.schemas.Schema) ReadOnlyTableProvider(org.apache.beam.sdk.extensions.sql.meta.provider.ReadOnlyTableProvider) JdbcConnection(org.apache.beam.sdk.extensions.sql.impl.JdbcConnection) Method(java.lang.reflect.Method) Matchers.containsString(org.hamcrest.Matchers.containsString) Row(org.apache.beam.sdk.values.Row) Test(org.junit.Test)

Aggregations

BeamRelNode (org.apache.beam.sdk.extensions.sql.impl.rel.BeamRelNode)246 Test (org.junit.Test)241 Row (org.apache.beam.sdk.values.Row)207 Schema (org.apache.beam.sdk.schemas.Schema)54 BeamIOSourceRel (org.apache.beam.sdk.extensions.sql.impl.rel.BeamIOSourceRel)38 ByteString (com.google.protobuf.ByteString)24 BeamCalcRel (org.apache.beam.sdk.extensions.sql.impl.rel.BeamCalcRel)22 Matchers.containsString (org.hamcrest.Matchers.containsString)13 Value (com.google.zetasql.Value)9 DateTime (org.joda.time.DateTime)8 BeamSqlEnv (org.apache.beam.sdk.extensions.sql.impl.BeamSqlEnv)7 BeamPushDownIOSourceRel (org.apache.beam.sdk.extensions.sql.impl.rel.BeamPushDownIOSourceRel)6 Calc (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.core.Calc)6 PipelineResult (org.apache.beam.sdk.PipelineResult)5 TestTableProvider (org.apache.beam.sdk.extensions.sql.meta.provider.test.TestTableProvider)5 DateTimeUtils.parseDateToValue (org.apache.beam.sdk.extensions.sql.zetasql.DateTimeUtils.parseDateToValue)5 DateTimeUtils.parseTimeToValue (org.apache.beam.sdk.extensions.sql.zetasql.DateTimeUtils.parseTimeToValue)5 DateTimeUtils.parseTimestampWithTZToValue (org.apache.beam.sdk.extensions.sql.zetasql.DateTimeUtils.parseTimestampWithTZToValue)5 Ignore (org.junit.Ignore)5 BeamAggregationRel (org.apache.beam.sdk.extensions.sql.impl.rel.BeamAggregationRel)4