Examples with TupleTag org.apache.beam.sdk.values.TupleTag used on opensource projects

Example 41 with TupleTag

use of org.apache.beam.sdk.values.TupleTag in project beam by apache.

the class BeamSqlDslUdfUdafTest method testSerializableFunctionUdf.

/**
 * Test UDF implementing {@link SerializableFunction}.
 */
@Test
public void testSerializableFunctionUdf() throws Exception {
    Schema resultType = Schema.builder().addInt32Field("f_int").addInt32Field("cubicvalue").build();
    Row row = Row.withSchema(resultType).addValues(2, 8).build();
    String sql = "SELECT f_int, cubic(f_int) as cubicvalue FROM PCOLLECTION WHERE f_int = 2";
    PCollection<Row> result = PCollectionTuple.of(new TupleTag<>("PCOLLECTION"), boundedInput1).apply("testUdf", SqlTransform.query(sql).registerUdf("cubic", new CubicIntegerFn()));
    PAssert.that(result).containsInAnyOrder(row);
    pipeline.run().waitUntilFinish();
}

Example 42 with TupleTag

use of org.apache.beam.sdk.values.TupleTag in project beam by apache.

the class BeamSqlDslProjectTest method testProjectUnknownField.

@Test
public void testProjectUnknownField() throws Exception {
    exceptions.expect(ParseException.class);
    exceptions.expectCause(hasMessage(containsString("Column 'f_int_na' not found in any table")));
    pipeline.enableAbandonedNodeEnforcement(false);
    String sql = "SELECT f_int_na FROM TABLE_A";
    PCollectionTuple.of(new TupleTag<>("TABLE_A"), boundedInput1).apply("testProjectUnknownField", SqlTransform.query(sql));
    pipeline.run();
}

Example 43 with TupleTag

use of org.apache.beam.sdk.values.TupleTag in project beam by apache.

the class BeamSqlDslProjectTest method runPartialFields.

private void runPartialFields(PCollection<Row> input) throws Exception {
    String sql = "SELECT f_int, f_long FROM TABLE_A";
    PCollection<Row> result = PCollectionTuple.of(new TupleTag<>("TABLE_A"), input).apply("testPartialFields", SqlTransform.query(sql));
    Schema resultType = Schema.builder().addInt32Field("f_int").addInt64Field("f_long").build();
    Row row = rowAtIndex(resultType, 0);
    PAssert.that(result).containsInAnyOrder(row);
    pipeline.run();
}

Example 44 with TupleTag

use of org.apache.beam.sdk.values.TupleTag in project beam by apache.

the class BeamSqlDslUdfUdafTest method testBeamSqlUdfWithDefaultParameters.

/**
 * Test UDF implementing {@link BeamSqlUdf} with default parameters.
 */
@Test
public void testBeamSqlUdfWithDefaultParameters() throws Exception {
    String sql = "SELECT f_int, substr(f_string) as sub_string FROM PCOLLECTION WHERE f_int = 2";
    PCollection<Row> result = PCollectionTuple.of(new TupleTag<>("PCOLLECTION"), boundedInput1).apply("testUdf", SqlTransform.query(sql).registerUdf("substr", UdfFnWithDefault.class));
    Schema subStrSchema = Schema.builder().addInt32Field("f_int").addStringField("sub_string").build();
    Row subStrRow = Row.withSchema(subStrSchema).addValues(2, "s").build();
    PAssert.that(result).containsInAnyOrder(subStrRow);
    pipeline.run().waitUntilFinish();
}

Example 45 with TupleTag

use of org.apache.beam.sdk.values.TupleTag in project beam by apache.

the class PCollectionOutputTagVisitor method visitValue.

@Override
public void visitValue(PValue value, Node producer) {
    for (Entry<ProjectionProducer<PTransform<?, ?>>, Map<PCollection<?>, FieldAccessDescriptor>> entry : pCollFieldAccess.entrySet()) {
        FieldAccessDescriptor fieldAccess = entry.getValue().get(value);
        if (fieldAccess == null) {
            continue;
        }
        BiMap<PCollection<?>, TupleTag<?>> outputs = ImmutableBiMap.copyOf(producer.getOutputs()).inverse();
        TupleTag<?> tag = outputs.get(value);
        Preconditions.checkArgumentNotNull(tag, "PCollection %s not found in outputs of producer %s", value, producer);
        ImmutableMap.Builder<TupleTag<?>, FieldAccessDescriptor> tagEntryBuilder = tagFieldAccess.build().get(entry.getKey());
        if (tagEntryBuilder == null) {
            tagEntryBuilder = ImmutableMap.builder();
            tagFieldAccess.put(entry.getKey(), tagEntryBuilder);
        }
        tagEntryBuilder.put(tag, fieldAccess);
    }
}