001/* 002 * Copyright (c) 2007-2017 Xplenty, Inc. All Rights Reserved. 003 * 004 * Project and contact information: http://www.cascading.org/ 005 * 006 * This file is part of the Cascading project. 007 * 008 * Licensed under the Apache License, Version 2.0 (the "License"); 009 * you may not use this file except in compliance with the License. 010 * You may obtain a copy of the License at 011 * 012 * http://www.apache.org/licenses/LICENSE-2.0 013 * 014 * Unless required by applicable law or agreed to in writing, software 015 * distributed under the License is distributed on an "AS IS" BASIS, 016 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 017 * See the License for the specific language governing permissions and 018 * limitations under the License. 019 */ 020 021package cascading.flow.planner.iso.subgraph.partitioner; 022 023import java.util.Collection; 024import java.util.LinkedHashMap; 025import java.util.Map; 026 027import cascading.flow.FlowElement; 028import cascading.flow.planner.PlannerContext; 029import cascading.flow.planner.graph.ElementGraph; 030import cascading.flow.planner.iso.ElementAnnotation; 031import cascading.flow.planner.iso.expression.ExpressionGraph; 032import cascading.flow.planner.iso.subgraph.GraphPartitioner; 033import cascading.flow.planner.iso.subgraph.Partitions; 034import cascading.flow.planner.iso.subgraph.SubGraphIterator; 035import cascading.flow.planner.iso.subgraph.iterator.ExpressionSubGraphIterator; 036import cascading.util.EnumMultiMap; 037 038/** 039 * 040 */ 041public class ExpressionGraphPartitioner extends GraphPartitioner 042 { 043 protected ExpressionGraph contractionGraph; 044 protected ExpressionGraph expressionGraph; 045 protected ElementAnnotation[] annotations = new ElementAnnotation[ 0 ]; 046 047 public ExpressionGraphPartitioner( ExpressionGraph contractionGraph, ExpressionGraph expressionGraph, ElementAnnotation... annotations ) 048 { 049 this.contractionGraph = contractionGraph; 050 this.expressionGraph = expressionGraph; 051 this.annotations = annotations; 052 } 053 054 public ExpressionGraph getContractionGraph() 055 { 056 return contractionGraph; 057 } 058 059 public ExpressionGraph getExpressionGraph() 060 { 061 return expressionGraph; 062 } 063 064 public ElementAnnotation[] getAnnotations() 065 { 066 return annotations; 067 } 068 069 public void setAnnotations( ElementAnnotation[] annotations ) 070 { 071 this.annotations = annotations; 072 } 073 074 @Override 075 public Partitions partition( PlannerContext plannerContext, ElementGraph elementGraph, Collection<FlowElement> excludes ) 076 { 077 Map<ElementGraph, EnumMultiMap> annotatedSubGraphs = new LinkedHashMap<>(); 078 079 ExpressionSubGraphIterator expressionIterator = new ExpressionSubGraphIterator( plannerContext, contractionGraph, expressionGraph, elementGraph, excludes ); 080 081 SubGraphIterator stepIterator = wrapIterator( expressionIterator ); 082 083 while( stepIterator.hasNext() ) 084 { 085 ElementGraph next = stepIterator.next(); 086 EnumMultiMap annotationMap = stepIterator.getAnnotationMap( annotations ); 087 088 annotatedSubGraphs.put( next, annotationMap ); 089 } 090 091 return new Partitions( this, elementGraph, expressionIterator.getContractedGraph(), expressionIterator.getMatches(), annotatedSubGraphs ); 092 } 093 094 protected SubGraphIterator wrapIterator( ExpressionSubGraphIterator expressionIterator ) 095 { 096 return expressionIterator; 097 } 098 }