Search in sources :

Example 76 with Table

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.schema.Table in project beam by apache.

the class DataStoreReadWriteIT method testDataStoreV1SqlWriteRead.

@Test
public void testDataStoreV1SqlWriteRead() {
    BeamSqlEnv sqlEnv = BeamSqlEnv.inMemory(new DataStoreV1TableProvider());
    String projectId = options.getProject();
    String createTableStatement = "CREATE EXTERNAL TABLE TEST( \n" + "   `__key__` VARBINARY, \n" + "   `content` VARCHAR \n" + ") \n" + "TYPE 'datastoreV1' \n" + "LOCATION '" + projectId + "/" + KIND + "'";
    sqlEnv.executeDdl(createTableStatement);
    Key ancestor = makeKey(KIND, UUID.randomUUID().toString()).build();
    Key itemKey = makeKey(ancestor, KIND, UUID.randomUUID().toString()).build();
    String insertStatement = "INSERT INTO TEST VALUES ( \n" + keyToSqlByteString(itemKey) + ", \n" + "'2000' \n" + ")";
    BeamSqlRelUtils.toPCollection(writePipeline, sqlEnv.parseQuery(insertStatement));
    writePipeline.run().waitUntilFinish();
    String selectTableStatement = "SELECT * FROM TEST";
    PCollection<Row> output = BeamSqlRelUtils.toPCollection(readPipeline, sqlEnv.parseQuery(selectTableStatement));
    assertThat(output.getSchema(), equalTo(SOURCE_SCHEMA));
    PipelineResult.State state = readPipeline.run().waitUntilFinish(Duration.standardMinutes(5));
    assertThat(state, equalTo(State.DONE));
}
Also used : State(org.apache.beam.sdk.PipelineResult.State) PipelineResult(org.apache.beam.sdk.PipelineResult) BeamSqlEnv(org.apache.beam.sdk.extensions.sql.impl.BeamSqlEnv) ByteString(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.avatica.util.ByteString) Row(org.apache.beam.sdk.values.Row) EntityToRow(org.apache.beam.sdk.io.gcp.datastore.EntityToRow) Key(com.google.datastore.v1.Key) DatastoreHelper.makeKey(com.google.datastore.v1.client.DatastoreHelper.makeKey) Test(org.junit.Test)

Example 77 with Table

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.schema.Table in project beam by apache.

the class BeamTableFunctionScanRule method convert.

@Override
public RelNode convert(RelNode relNode) {
    TableFunctionScan tableFunctionScan = (TableFunctionScan) relNode;
    // only support one input for table function scan.
    List<RelNode> inputs = new ArrayList<>();
    checkArgument(relNode.getInputs().size() == 1, "Wrong number of inputs for %s, expected 1 input but received: %s", BeamTableFunctionScanRel.class.getSimpleName(), relNode.getInputs().size());
    inputs.add(convert(relNode.getInput(0), relNode.getInput(0).getTraitSet().replace(BeamLogicalConvention.INSTANCE)));
    return new BeamTableFunctionScanRel(tableFunctionScan.getCluster(), tableFunctionScan.getTraitSet().replace(BeamLogicalConvention.INSTANCE), inputs, tableFunctionScan.getCall(), null, tableFunctionScan.getCall().getType(), Collections.EMPTY_SET);
}
Also used : BeamTableFunctionScanRel(org.apache.beam.sdk.extensions.sql.impl.rel.BeamTableFunctionScanRel) LogicalTableFunctionScan(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.logical.LogicalTableFunctionScan) TableFunctionScan(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.core.TableFunctionScan) RelNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelNode) ArrayList(java.util.ArrayList)

Example 78 with Table

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.schema.Table in project beam by apache.

the class ExpressionConverter method convertTableValuedFunction.

/**
 * Convert a TableValuedFunction in ZetaSQL to a RexCall in Calcite.
 */
public RexCall convertTableValuedFunction(RelNode input, TableValuedFunction tvf, List<ResolvedNodes.ResolvedFunctionArgument> argumentList, List<ResolvedColumn> inputTableColumns) {
    ResolvedColumn wmCol;
    // Handle builtin windowing TVF.
    switch(tvf.getName()) {
        case TVFStreamingUtils.FIXED_WINDOW_TVF:
            // TUMBLE tvf's second argument is descriptor.
            wmCol = extractWatermarkColumnFromDescriptor(argumentList.get(1).getDescriptorArg());
            return (RexCall) rexBuilder().makeCall(new SqlWindowTableFunction(SqlKind.TUMBLE.name()), convertRelNodeToRexRangeRef(input), convertResolvedColumnToRexInputRef(wmCol, inputTableColumns), convertIntervalToRexIntervalLiteral((ResolvedLiteral) argumentList.get(2).getExpr()));
        case TVFStreamingUtils.SLIDING_WINDOW_TVF:
            // HOP tvf's second argument is descriptor.
            wmCol = extractWatermarkColumnFromDescriptor(argumentList.get(1).getDescriptorArg());
            return (RexCall) rexBuilder().makeCall(new SqlWindowTableFunction(SqlKind.HOP.name()), convertRelNodeToRexRangeRef(input), convertResolvedColumnToRexInputRef(wmCol, inputTableColumns), convertIntervalToRexIntervalLiteral((ResolvedLiteral) argumentList.get(2).getExpr()), convertIntervalToRexIntervalLiteral((ResolvedLiteral) argumentList.get(3).getExpr()));
        case TVFStreamingUtils.SESSION_WINDOW_TVF:
            // SESSION tvf's second argument is descriptor.
            wmCol = extractWatermarkColumnFromDescriptor(argumentList.get(1).getDescriptorArg());
            // SESSION tvf's third argument is descriptor.
            List<ResolvedColumn> keyCol = extractSessionKeyColumnFromDescriptor(argumentList.get(2).getDescriptorArg());
            List<RexNode> operands = new ArrayList<>();
            operands.add(convertRelNodeToRexRangeRef(input));
            operands.add(convertResolvedColumnToRexInputRef(wmCol, inputTableColumns));
            operands.add(convertIntervalToRexIntervalLiteral((ResolvedLiteral) argumentList.get(3).getExpr()));
            operands.addAll(convertResolvedColumnsToRexInputRef(keyCol, inputTableColumns));
            return (RexCall) rexBuilder().makeCall(new SqlWindowTableFunction(SqlKind.SESSION.name()), operands);
    }
    if (tvf instanceof FixedOutputSchemaTVF) {
        FixedOutputSchemaTVF fixedOutputSchemaTVF = (FixedOutputSchemaTVF) tvf;
        return (RexCall) rexBuilder().makeCall(new ZetaSqlUserDefinedSQLNativeTableValuedFunction(new SqlIdentifier(tvf.getName(), SqlParserPos.ZERO), opBinding -> {
            List<RelDataTypeField> relDataTypeFields = convertTVFRelationColumnsToRelDataTypeFields(fixedOutputSchemaTVF.getOutputSchema().getColumns());
            return new RelRecordType(relDataTypeFields);
        }, null, null, null, null));
    }
    throw new UnsupportedOperationException("Does not support table-valued function: " + tvf.getName());
}
Also used : FixedOutputSchemaTVF(com.google.zetasql.TableValuedFunction.FixedOutputSchemaTVF) ZetaSqlUserDefinedSQLNativeTableValuedFunction(org.apache.beam.sdk.extensions.sql.impl.ZetaSqlUserDefinedSQLNativeTableValuedFunction) ArrayList(java.util.ArrayList) SqlIdentifier(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlIdentifier) ResolvedColumn(com.google.zetasql.resolvedast.ResolvedColumn) RelRecordType(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelRecordType) RexCall(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexCall) RelDataTypeField(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataTypeField) ResolvedLiteral(com.google.zetasql.resolvedast.ResolvedNodes.ResolvedLiteral) RexNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode)

Example 79 with Table

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.schema.Table in project beam by apache.

the class ZetaSQLPlannerImpl method rel.

public RelRoot rel(String sql, QueryParameters params) {
    RelOptCluster cluster = RelOptCluster.create(planner, new RexBuilder(typeFactory));
    AnalyzerOptions options = SqlAnalyzer.getAnalyzerOptions(params, defaultTimezone);
    BeamZetaSqlCatalog catalog = BeamZetaSqlCatalog.create(defaultSchemaPlus, (JavaTypeFactory) cluster.getTypeFactory(), options);
    // Set up table providers that need to be pre-registered
    SqlAnalyzer analyzer = new SqlAnalyzer();
    List<List<String>> tables = analyzer.extractTableNames(sql, options);
    TableResolution.registerTables(this.defaultSchemaPlus, tables);
    QueryTrait trait = new QueryTrait();
    catalog.addTables(tables, trait);
    ResolvedQueryStmt statement = analyzer.analyzeQuery(sql, options, catalog);
    ExpressionConverter expressionConverter = new ExpressionConverter(cluster, params, catalog.getUserFunctionDefinitions());
    ConversionContext context = ConversionContext.of(config, expressionConverter, cluster, trait);
    RelNode convertedNode = QueryStatementConverter.convertRootQuery(context, statement);
    return RelRoot.of(convertedNode, SqlKind.ALL);
}
Also used : RelOptCluster(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.plan.RelOptCluster) ConversionContext(org.apache.beam.sdk.extensions.sql.zetasql.translation.ConversionContext) RelNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelNode) RexBuilder(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexBuilder) ExpressionConverter(org.apache.beam.sdk.extensions.sql.zetasql.translation.ExpressionConverter) List(java.util.List) ImmutableList(org.apache.beam.vendor.calcite.v1_28_0.com.google.common.collect.ImmutableList) ResolvedQueryStmt(com.google.zetasql.resolvedast.ResolvedNodes.ResolvedQueryStmt) AnalyzerOptions(com.google.zetasql.AnalyzerOptions)

Example 80 with Table

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.schema.Table in project beam by apache.

the class PubsubTableProviderIT method testSQLSelectsPayloadContent.

@Test
public void testSQLSelectsPayloadContent() throws Exception {
    String createTableString = String.format("CREATE EXTERNAL TABLE message (\n" + "event_timestamp TIMESTAMP, \n" + "attributes MAP<VARCHAR, VARCHAR>, \n" + "payload ROW< \n" + "             id INTEGER, \n" + "             name VARCHAR \n" + "           > \n" + ") \n" + "TYPE '%s' \n" + "LOCATION '%s' \n" + "TBLPROPERTIES '{ " + "%s" + "\"protoClass\" : \"%s\", " + "\"timestampAttributeKey\" : \"ts\" }'", tableProvider.getTableType(), eventsTopic.topicPath(), payloadFormatParam(), PayloadMessages.SimpleMessage.class.getName());
    String queryString = "SELECT message.payload.id, message.payload.name from message";
    // Initialize SQL environment and create the pubsub table
    BeamSqlEnv sqlEnv = BeamSqlEnv.inMemory(new PubsubTableProvider());
    sqlEnv.executeDdl(createTableString);
    // Apply the PTransform to query the pubsub topic
    PCollection<Row> queryOutput = query(sqlEnv, pipeline, queryString);
    // Observe the query results and send success signal after seeing the expected messages
    queryOutput.apply("waitForSuccess", resultSignal.signalSuccessWhen(SchemaCoder.of(PAYLOAD_SCHEMA), observedRows -> observedRows.equals(ImmutableSet.of(row(PAYLOAD_SCHEMA, 3, "foo"), row(PAYLOAD_SCHEMA, 5, "bar"), row(PAYLOAD_SCHEMA, 7, "baz")))));
    // Start the pipeline
    pipeline.run();
    // Block until a subscription for this topic exists
    eventsTopic.assertSubscriptionEventuallyCreated(pipeline.getOptions().as(GcpOptions.class).getProject(), Duration.standardMinutes(5));
    // Start publishing the messages when main pipeline is started and signaling topic is ready
    eventsTopic.publish(ImmutableList.of(objectsProvider.messageIdName(ts(1), 3, "foo"), objectsProvider.messageIdName(ts(2), 5, "bar"), objectsProvider.messageIdName(ts(3), 7, "baz")));
    // Poll the signaling topic for success message
    resultSignal.waitForSuccess(timeout);
}
Also used : Arrays(java.util.Arrays) LoggerFactory(org.slf4j.LoggerFactory) TimeoutException(java.util.concurrent.TimeoutException) PubsubMessage(org.apache.beam.sdk.io.gcp.pubsub.PubsubMessage) Future(java.util.concurrent.Future) TestPubsub(org.apache.beam.sdk.io.gcp.pubsub.TestPubsub) ResultSet(java.sql.ResultSet) Map(java.util.Map) TestPubsubSignal(org.apache.beam.sdk.io.gcp.pubsub.TestPubsubSignal) Parameterized(org.junit.runners.Parameterized) ImmutableMap(org.apache.beam.vendor.calcite.v1_28_0.com.google.common.collect.ImmutableMap) GcpOptions(org.apache.beam.sdk.extensions.gcp.options.GcpOptions) Matchers.allOf(org.hamcrest.Matchers.allOf) Collection(java.util.Collection) SchemaCoder(org.apache.beam.sdk.schemas.SchemaCoder) Set(java.util.Set) FieldType(org.apache.beam.sdk.schemas.Schema.FieldType) SchemaIOTableProviderWrapper(org.apache.beam.sdk.extensions.sql.meta.provider.SchemaIOTableProviderWrapper) Collectors(java.util.stream.Collectors) StandardCharsets(java.nio.charset.StandardCharsets) Executors(java.util.concurrent.Executors) ImmutableSet(org.apache.beam.vendor.calcite.v1_28_0.com.google.common.collect.ImmutableSet) Serializable(java.io.Serializable) List(java.util.List) Matchers.equalTo(org.hamcrest.Matchers.equalTo) JdbcDriver(org.apache.beam.sdk.extensions.sql.impl.JdbcDriver) ReflectHelpers(org.apache.beam.sdk.util.common.ReflectHelpers) ImmutableList(org.apache.beam.vendor.calcite.v1_28_0.com.google.common.collect.ImmutableList) GenericRecordBuilder(org.apache.avro.generic.GenericRecordBuilder) JsonMatcher.jsonBytesLike(org.apache.beam.sdk.testing.JsonMatcher.jsonBytesLike) ByteArrayOutputStream(java.io.ByteArrayOutputStream) InMemoryMetaStore(org.apache.beam.sdk.extensions.sql.meta.store.InMemoryMetaStore) Duration(org.joda.time.Duration) RunWith(org.junit.runner.RunWith) Parameters(org.junit.runners.Parameterized.Parameters) HashMap(java.util.HashMap) Callable(java.util.concurrent.Callable) Matchers.hasProperty(org.hamcrest.Matchers.hasProperty) BeamSqlRelUtils(org.apache.beam.sdk.extensions.sql.impl.rel.BeamSqlRelUtils) TestPipeline(org.apache.beam.sdk.testing.TestPipeline) MatcherAssert.assertThat(org.hamcrest.MatcherAssert.assertThat) BeamSqlEnv(org.apache.beam.sdk.extensions.sql.impl.BeamSqlEnv) Row(org.apache.beam.sdk.values.Row) PipelineOptions(org.apache.beam.sdk.options.PipelineOptions) ExecutorService(java.util.concurrent.ExecutorService) AvroUtils(org.apache.beam.sdk.schemas.utils.AvroUtils) Matchers.hasEntry(org.hamcrest.Matchers.hasEntry) GenericRecord(org.apache.avro.generic.GenericRecord) Logger(org.slf4j.Logger) UTF_8(java.nio.charset.StandardCharsets.UTF_8) Parameter(org.junit.runners.Parameterized.Parameter) ObjectMapper(com.fasterxml.jackson.databind.ObjectMapper) JdbcConnection(org.apache.beam.sdk.extensions.sql.impl.JdbcConnection) JsonProcessingException(com.fasterxml.jackson.core.JsonProcessingException) IOException(java.io.IOException) Test(org.junit.Test) PCollection(org.apache.beam.sdk.values.PCollection) AvroCoder(org.apache.beam.sdk.coders.AvroCoder) Schema(org.apache.beam.sdk.schemas.Schema) ExecutionException(java.util.concurrent.ExecutionException) TimeUnit(java.util.concurrent.TimeUnit) PayloadMessages(org.apache.beam.sdk.extensions.protobuf.PayloadMessages) Rule(org.junit.Rule) Ignore(org.junit.Ignore) Matcher(org.hamcrest.Matcher) Instant(org.joda.time.Instant) Statement(java.sql.Statement) TableProvider(org.apache.beam.sdk.extensions.sql.meta.provider.TableProvider) CalciteConnection(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.jdbc.CalciteConnection) BeamSqlEnv(org.apache.beam.sdk.extensions.sql.impl.BeamSqlEnv) Row(org.apache.beam.sdk.values.Row) Test(org.junit.Test)

Aggregations

Table (org.apache.calcite.schema.Table)104 Test (org.junit.Test)43 SqlStdOperatorTable (org.apache.calcite.sql.fun.SqlStdOperatorTable)38 RelOptTable (org.apache.calcite.plan.RelOptTable)33 RelDataType (org.apache.calcite.rel.type.RelDataType)22 SqlOperatorTable (org.apache.calcite.sql.SqlOperatorTable)22 SchemaPlus (org.apache.calcite.schema.SchemaPlus)20 ArrayList (java.util.ArrayList)19 List (java.util.List)19 ProjectableFilterableTable (org.apache.calcite.schema.ProjectableFilterableTable)17 ScannableTable (org.apache.calcite.schema.ScannableTable)17 JavaTypeFactoryImpl (org.apache.calcite.jdbc.JavaTypeFactoryImpl)16 FilterableTable (org.apache.calcite.schema.FilterableTable)15 AbstractTable (org.apache.calcite.schema.impl.AbstractTable)15 ResultSet (java.sql.ResultSet)14 Schema (org.apache.beam.sdk.schemas.Schema)14 Row (org.apache.beam.sdk.values.Row)14 StreamableTable (org.apache.calcite.schema.StreamableTable)14 CalciteConnection (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.jdbc.CalciteConnection)13 Map (java.util.Map)12