Examples with ProjectionProducer org.apache.beam.sdk.schemas.ProjectionProducer used on opensource projects

Example 6 with ProjectionProducer

use of org.apache.beam.sdk.schemas.ProjectionProducer in project beam by apache.

the class ProjectionProducerVisitorTest method testNestedPushdownProducers_returnsOnlyOutermostPushdown.

@Test
public void testNestedPushdownProducers_returnsOnlyOutermostPushdown() {
    Pipeline p = Pipeline.create();
    PTransform<PBegin, PCollection<Row>> source = new CompositeTransformWithPushdownOutside();
    PCollection<Row> output = p.apply(source);
    Map<PCollection<?>, FieldAccessDescriptor> pCollectionFieldAccess = ImmutableMap.of(output, FieldAccessDescriptor.withFieldNames("field1", "field2"));
    ProjectionProducerVisitor visitor = new ProjectionProducerVisitor(pCollectionFieldAccess);
    p.traverseTopologically(visitor);
    Map<ProjectionProducer<PTransform<?, ?>>, Map<PCollection<?>, FieldAccessDescriptor>> pushdownOpportunities = visitor.getPushdownOpportunities();
    Assert.assertEquals(1, pushdownOpportunities.size());
    Map<PCollection<?>, FieldAccessDescriptor> opportunitiesForSource = pushdownOpportunities.get(source);
    Assert.assertNotNull(opportunitiesForSource);
    Assert.assertEquals(1, opportunitiesForSource.size());
    FieldAccessDescriptor fieldAccessDescriptor = opportunitiesForSource.get(output);
    Assert.assertNotNull(fieldAccessDescriptor);
    Assert.assertFalse(fieldAccessDescriptor.getAllFields());
    assertThat(fieldAccessDescriptor.fieldNamesAccessed(), containsInAnyOrder("field1", "field2"));
}

Example 7 with ProjectionProducer

use of org.apache.beam.sdk.schemas.ProjectionProducer in project beam by apache.

the class ProjectionProducerVisitorTest method testSimplePushdownProducer_returnsOnePushdown.

@Test
public void testSimplePushdownProducer_returnsOnePushdown() {
    Pipeline p = Pipeline.create();
    PTransform<PBegin, PCollection<Row>> source = new SimpleSourceWithPushdown();
    PCollection<Row> output = p.apply(source);
    Map<PCollection<?>, FieldAccessDescriptor> pCollectionFieldAccess = ImmutableMap.of(output, FieldAccessDescriptor.withFieldNames("field1", "field2"));
    ProjectionProducerVisitor visitor = new ProjectionProducerVisitor(pCollectionFieldAccess);
    p.traverseTopologically(visitor);
    Map<ProjectionProducer<PTransform<?, ?>>, Map<PCollection<?>, FieldAccessDescriptor>> pushdownOpportunities = visitor.getPushdownOpportunities();
    Assert.assertEquals(1, pushdownOpportunities.size());
    Map<PCollection<?>, FieldAccessDescriptor> opportunitiesForSource = pushdownOpportunities.get(source);
    Assert.assertNotNull(opportunitiesForSource);
    Assert.assertEquals(1, opportunitiesForSource.size());
    FieldAccessDescriptor fieldAccessDescriptor = opportunitiesForSource.get(output);
    Assert.assertNotNull(fieldAccessDescriptor);
    Assert.assertFalse(fieldAccessDescriptor.getAllFields());
    assertThat(fieldAccessDescriptor.fieldNamesAccessed(), containsInAnyOrder("field1", "field2"));
}

Example 8 with ProjectionProducer

use of org.apache.beam.sdk.schemas.ProjectionProducer in project beam by apache.

the class ProjectionProducerVisitorTest method testProjectionProducerInsideNonProducer_returnsInnerPushdown.

@Test
public void testProjectionProducerInsideNonProducer_returnsInnerPushdown() {
    Pipeline p = Pipeline.create();
    CompositeTransformWithPushdownInside source = new CompositeTransformWithPushdownInside();
    PCollection<Row> output = p.apply(source);
    Map<PCollection<?>, FieldAccessDescriptor> pCollectionFieldAccess = ImmutableMap.of(output, FieldAccessDescriptor.withFieldNames("field1", "field2"));
    ProjectionProducerVisitor visitor = new ProjectionProducerVisitor(pCollectionFieldAccess);
    p.traverseTopologically(visitor);
    Map<ProjectionProducer<PTransform<?, ?>>, Map<PCollection<?>, FieldAccessDescriptor>> pushdownOpportunities = visitor.getPushdownOpportunities();
    Assert.assertEquals(1, pushdownOpportunities.size());
    Map<PCollection<?>, FieldAccessDescriptor> opportunitiesForSource = pushdownOpportunities.get(source.innerT);
    Assert.assertNotNull(opportunitiesForSource);
    Assert.assertEquals(1, opportunitiesForSource.size());
    FieldAccessDescriptor fieldAccessDescriptor = opportunitiesForSource.get(output);
    Assert.assertNotNull(fieldAccessDescriptor);
    Assert.assertFalse(fieldAccessDescriptor.getAllFields());
    assertThat(fieldAccessDescriptor.fieldNamesAccessed(), containsInAnyOrder("field1", "field2"));
}

Example 9 with ProjectionProducer

use of org.apache.beam.sdk.schemas.ProjectionProducer in project beam by apache.

the class ProjectionProducerVisitorTest method testFieldAccessAllFields_returnsNoPushdown.

@Test
public void testFieldAccessAllFields_returnsNoPushdown() {
    Pipeline p = Pipeline.create();
    PCollection<Row> output = p.apply(new SimpleSource());
    Map<PCollection<?>, FieldAccessDescriptor> pCollectionFieldAccess = ImmutableMap.of(output, FieldAccessDescriptor.withAllFields());
    ProjectionProducerVisitor visitor = new ProjectionProducerVisitor(pCollectionFieldAccess);
    p.traverseTopologically(visitor);
    Map<ProjectionProducer<PTransform<?, ?>>, Map<PCollection<?>, FieldAccessDescriptor>> pushdownOpportunities = visitor.getPushdownOpportunities();
    Assert.assertTrue(pushdownOpportunities.isEmpty());
}