use of org.apache.beam.sdk.values.Row 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));
}
use of org.apache.beam.sdk.values.Row in project beam by apache.
the class ZetaSqlJavaUdfTest method testSqlTransformRegisterUdf.
@Test
public void testSqlTransformRegisterUdf() {
String sql = "SELECT increment(0);";
PCollection<Row> stream = pipeline.apply(SqlTransform.query(sql).withQueryPlannerClass(ZetaSQLQueryPlanner.class).registerUdf("increment", IncrementFn.class));
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));
}
use of org.apache.beam.sdk.values.Row 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));
}
use of org.apache.beam.sdk.values.Row in project beam by apache.
the class ZetaSqlJavaUdfTest method testNullaryJavaUdf.
@Test
public void testNullaryJavaUdf() {
String sql = String.format("CREATE FUNCTION helloWorld() RETURNS STRING LANGUAGE java " + "OPTIONS (path='%s'); " + "SELECT helloWorld();", jarPath);
ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config);
BeamRelNode beamRelNode = zetaSQLQueryPlanner.convertToBeamRel(sql);
PCollection<Row> stream = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
Schema singleField = Schema.builder().addStringField("field1").build();
PAssert.that(stream).containsInAnyOrder(Row.withSchema(singleField).addValues("Hello world!").build());
pipeline.run().waitUntilFinish(Duration.standardMinutes(PIPELINE_EXECUTION_WAITTIME_MINUTES));
}
use of org.apache.beam.sdk.values.Row in project beam by apache.
the class ZetaSqlJavaUdfTest method testProjectUdfAndBuiltin.
@Test
public void testProjectUdfAndBuiltin() {
String sql = String.format("CREATE FUNCTION matches(str STRING, regStr STRING) RETURNS BOOLEAN LANGUAGE java OPTIONS (path='%s'); " + "SELECT matches(\"a\", \"a\"), 'apple'='beta'", jarPath);
ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config);
BeamRelNode beamRelNode = zetaSQLQueryPlanner.convertToBeamRel(sql);
PCollection<Row> stream = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
Schema schema = Schema.builder().addBooleanField("field1").addBooleanField("field2").build();
PAssert.that(stream).containsInAnyOrder(Row.withSchema(schema).addValues(true, false).build());
pipeline.run().waitUntilFinish(Duration.standardMinutes(PIPELINE_EXECUTION_WAITTIME_MINUTES));
}
Aggregations