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

Example 11 with TableProvider

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

the class KafkaTableProviderIT method testFakeNested.

@Test
public void testFakeNested() throws InterruptedException {
    Assume.assumeFalse(topic.equals("csv_topic"));
    pipeline.getOptions().as(DirectOptions.class).setBlockOnRun(false);
    String createTableString = String.format("CREATE EXTERNAL TABLE kafka_table(\n" + "headers ARRAY<ROW<key VARCHAR, `values` ARRAY<VARBINARY>>>," + "payload ROW<" + "f_long BIGINT NOT NULL, \n" + "f_int INTEGER NOT NULL, \n" + "f_string VARCHAR NOT NULL \n" + ">" + ") \n" + "TYPE 'kafka' \n" + "LOCATION '%s'\n" + "TBLPROPERTIES '%s'", buildLocation(), objectsProvider.getKafkaPropertiesString());
    TableProvider tb = new KafkaTableProvider();
    BeamSqlEnv env = BeamSqlEnv.inMemory(tb);
    env.executeDdl(createTableString);
    PCollection<Row> queryOutput = BeamSqlRelUtils.toPCollection(pipeline, env.parseQuery("SELECT kafka_table.payload.f_long, kafka_table.payload.f_int, kafka_table.payload.f_string FROM kafka_table"));
    queryOutput.apply(ParDo.of(new FakeKvPair())).setCoder(KvCoder.of(StringUtf8Coder.of(), RowCoder.of(TEST_TABLE_SCHEMA))).apply("waitForSuccess", ParDo.of(new StreamAssertEqual(ImmutableSet.of(generateRow(0), generateRow(1), generateRow(2)))));
    queryOutput.apply(logRecords(""));
    pipeline.run();
    TimeUnit.SECONDS.sleep(4);
    produceSomeRecords(3);
    for (int i = 0; i < 200; i++) {
        if (FLAG.getOrDefault(pipeline.getOptions().getOptionsId(), false)) {
            return;
        }
        TimeUnit.MILLISECONDS.sleep(90);
    }
    Assert.fail();
}

Example 12 with TableProvider

use of org.apache.beam.sdk.extensions.sql.meta.provider.TableProvider 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 13 with TableProvider

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

the class PubsubTableProviderIT method connect.

@SuppressWarnings("unchecked")
private CalciteConnection connect(PipelineOptions options, TableProvider... tableProviders) {
    // HACK: PipelineOptions should expose a prominent method to do this reliably
    // The actual options are in the "options" field of the converted map
    Map<String, String> argsMap = ((Map<String, Object>) MAPPER.convertValue(pipeline.getOptions(), Map.class).get("options")).entrySet().stream().filter((entry) -> {
        if (entry.getValue() instanceof List) {
            if (!((List) entry.getValue()).isEmpty()) {
                throw new IllegalArgumentException("Cannot encode list arguments");
            }
            // We can encode empty lists, just omit them.
            return false;
        }
        return true;
    }).collect(Collectors.toMap(Map.Entry::getKey, entry -> toArg(entry.getValue())));
    InMemoryMetaStore inMemoryMetaStore = new InMemoryMetaStore();
    for (TableProvider tableProvider : tableProviders) {
        inMemoryMetaStore.registerProvider(tableProvider);
    }
    JdbcConnection connection = JdbcDriver.connect(inMemoryMetaStore, options);
    connection.setPipelineOptionsMap(argsMap);
    return connection;
}