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"));
}
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"));
}
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"));
}
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());
}
Aggregations