Examples with ReadOnlyTableProvider org.apache.beam.sdk.extensions.sql.meta.provider.ReadOnlyTableProvider used on opensource projects

Example 1 with ReadOnlyTableProvider

use of org.apache.beam.sdk.extensions.sql.meta.provider.ReadOnlyTableProvider 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));
}

Example 2 with ReadOnlyTableProvider

use of org.apache.beam.sdk.extensions.sql.meta.provider.ReadOnlyTableProvider in project beam by apache.

the class ZetaSqlJavaUdfTypeTest method setUp.

@Before
public void setUp() throws NoSuchMethodException {
    initialize();
    // Register test table.
    JdbcConnection jdbcConnection = JdbcDriver.connect(new ReadOnlyTableProvider("table_provider", ImmutableMap.of("table", table)), PipelineOptionsFactory.create());
    // Register UDFs.
    SchemaPlus schema = jdbcConnection.getCurrentSchemaPlus();
    schema.add("test_boolean", ScalarFunctionImpl.create(BooleanIdentityFn.class.getMethod("eval", Boolean.class)));
    schema.add("test_int64", ScalarFunctionImpl.create(Int64IdentityFn.class.getMethod("eval", Long.class)));
    schema.add("test_string", ScalarFunctionImpl.create(StringIdentityFn.class.getMethod("eval", String.class)));
    schema.add("test_bytes", ScalarFunctionImpl.create(BytesIdentityFn.class.getMethod("eval", byte[].class)));
    schema.add("test_float64", ScalarFunctionImpl.create(DoubleIdentityFn.class.getMethod("eval", Double.class)));
    schema.add("test_date", ScalarFunctionImpl.create(DateIdentityFn.class.getMethod("eval", Date.class)));
    schema.add("test_timestamp", ScalarFunctionImpl.create(TimestampIdentityFn.class.getMethod("eval", Timestamp.class)));
    schema.add("test_array", ScalarFunctionImpl.create(ListIdentityFn.class.getMethod("eval", List.class)));
    schema.add("test_numeric", ScalarFunctionImpl.create(BigDecimalIdentityFn.class.getMethod("eval", BigDecimal.class)));
    this.config = Frameworks.newConfigBuilder(config).defaultSchema(schema).build();
}

Example 3 with ReadOnlyTableProvider

use of org.apache.beam.sdk.extensions.sql.meta.provider.ReadOnlyTableProvider in project beam by apache.

the class JdbcDriverTest method testInternalConnect_bounded_limit.

@Test
public void testInternalConnect_bounded_limit() throws Exception {
    ReadOnlyTableProvider tableProvider = new ReadOnlyTableProvider("test", ImmutableMap.of("test", TestBoundedTable.of(Schema.FieldType.INT32, "id", Schema.FieldType.STRING, "name").addRows(1, "first").addRows(1, "second first").addRows(2, "second")));
    CalciteConnection connection = JdbcDriver.connect(tableProvider, PipelineOptionsFactory.create());
    Statement statement = connection.createStatement();
    ResultSet resultSet1 = statement.executeQuery("SELECT * FROM test LIMIT 5");
    assertTrue(resultSet1.next());
    assertTrue(resultSet1.next());
    assertTrue(resultSet1.next());
    assertFalse(resultSet1.next());
    assertFalse(resultSet1.next());
    ResultSet resultSet2 = statement.executeQuery("SELECT * FROM test LIMIT 1");
    assertTrue(resultSet2.next());
    assertFalse(resultSet2.next());
    ResultSet resultSet3 = statement.executeQuery("SELECT * FROM test LIMIT 2");
    assertTrue(resultSet3.next());
    assertTrue(resultSet3.next());
    assertFalse(resultSet3.next());
    ResultSet resultSet4 = statement.executeQuery("SELECT * FROM test LIMIT 3");
    assertTrue(resultSet4.next());
    assertTrue(resultSet4.next());
    assertTrue(resultSet4.next());
    assertFalse(resultSet4.next());
}

Example 4 with ReadOnlyTableProvider

use of org.apache.beam.sdk.extensions.sql.meta.provider.ReadOnlyTableProvider in project beam by apache.

the class SqlTransform method expand.

@Override
public PCollection<Row> expand(PInput input) {
    TableProvider inputTableProvider = new ReadOnlyTableProvider(PCOLLECTION_NAME, toTableMap(input));
    InMemoryMetaStore metaTableProvider = new InMemoryMetaStore();
    metaTableProvider.registerProvider(inputTableProvider);
    BeamSqlEnvBuilder sqlEnvBuilder = BeamSqlEnv.builder(metaTableProvider);
    // TODO: validate duplicate functions.
    registerFunctions(sqlEnvBuilder);
    // the same names are reused.
    if (autoLoading()) {
        sqlEnvBuilder.autoLoadUserDefinedFunctions();
        ServiceLoader.load(TableProvider.class).forEach(metaTableProvider::registerProvider);
    }
    tableProviderMap().forEach(sqlEnvBuilder::addSchema);
    @Nullable final String defaultTableProvider = defaultTableProvider();
    if (defaultTableProvider != null) {
        sqlEnvBuilder.setCurrentSchema(defaultTableProvider);
    }
    sqlEnvBuilder.setQueryPlannerClassName(MoreObjects.firstNonNull(queryPlannerClassName(), input.getPipeline().getOptions().as(BeamSqlPipelineOptions.class).getPlannerName()));
    sqlEnvBuilder.setPipelineOptions(input.getPipeline().getOptions());
    BeamSqlEnv sqlEnv = sqlEnvBuilder.build();
    ddlStrings().forEach(sqlEnv::executeDdl);
    return BeamSqlRelUtils.toPCollection(input.getPipeline(), sqlEnv.parseQuery(queryString(), queryParameters()), errorsTransformer());
}

Example 5 with ReadOnlyTableProvider

use of org.apache.beam.sdk.extensions.sql.meta.provider.ReadOnlyTableProvider in project beam by apache.

the class JdbcDriverTest method testInternalConnect_unbounded_limit.

@Test
public void testInternalConnect_unbounded_limit() throws Exception {
    ReadOnlyTableProvider tableProvider = new ReadOnlyTableProvider("test", ImmutableMap.of("test", TestUnboundedTable.of(Schema.FieldType.INT32, "order_id", Schema.FieldType.INT32, "site_id", Schema.FieldType.INT32, "price", Schema.FieldType.DATETIME, "order_time").timestampColumnIndex(3).addRows(Duration.ZERO, 1, 1, 1, FIRST_DATE, 1, 2, 6, FIRST_DATE)));
    CalciteConnection connection = JdbcDriver.connect(tableProvider, PipelineOptionsFactory.create());
    Statement statement = connection.createStatement();
    ResultSet resultSet1 = statement.executeQuery("SELECT * FROM test LIMIT 1");
    assertTrue(resultSet1.next());
    assertFalse(resultSet1.next());
    ResultSet resultSet2 = statement.executeQuery("SELECT * FROM test LIMIT 2");
    assertTrue(resultSet2.next());
    assertTrue(resultSet2.next());
    assertFalse(resultSet2.next());
}