    001/*
    002 * Copyright (c) 2016 Chris K Wensel <chris@wensel.net>. All Rights Reserved.
    003 * Copyright (c) 2007-2017 Xplenty, Inc. All Rights Reserved.
    004 *
    005 * Project and contact information: http://www.cascading.org/
    006 *
    007 * This file is part of the Cascading project.
    008 *
    009 * Licensed under the Apache License, Version 2.0 (the "License");
    010 * you may not use this file except in compliance with the License.
    011 * You may obtain a copy of the License at
    012 *
    013 *     http://www.apache.org/licenses/LICENSE-2.0
    014 *
    015 * Unless required by applicable law or agreed to in writing, software
    016 * distributed under the License is distributed on an "AS IS" BASIS,
    017 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    018 * See the License for the specific language governing permissions and
    019 * limitations under the License.
    020 */
    021
    022package cascading.flow.planner.rule;
    023
    024import java.util.Collection;
    025
    026import cascading.flow.FlowElement;
    027import cascading.flow.planner.PlannerContext;
    028import cascading.flow.planner.graph.ElementGraph;
    029import cascading.flow.planner.iso.subgraph.GraphPartitioner;
    030import cascading.flow.planner.iso.subgraph.Partitions;
    031import cascading.flow.planner.rule.util.LogLevel;
    032
    033import static cascading.flow.planner.rule.util.RuleLogUtil.enableLogging;
    034import static cascading.flow.planner.rule.util.RuleLogUtil.restoreLogging;
    035
    036/**
    037 * The RulePartitioner class is responsible for partitioning an element graph into smaller sub-graphs.
    038 * <p/>
    039 * It may also re-partition a given graph, in place replacing it with its children, if any.
    040 */
    041public abstract class RulePartitioner implements Rule
    042  {
    043  public abstract Enum[] getAnnotationExcludes();
    044
    045  public enum PartitionSource
    046    {
    047      /**
    048       * Partition the parent into children.
    049       */
    050      PartitionParent,
    051
    052      /**
    053       * Partition a given child into more children, removing the original child.
    054       */
    055      PartitionCurrent
    056    }
    057
    058  protected LogLevel logLevel;
    059  protected PlanPhase phase;
    060  protected PartitionSource partitionSource = PartitionSource.PartitionParent;
    061  protected GraphPartitioner graphPartitioner;
    062
    063  public RulePartitioner( PlanPhase phase, PartitionSource partitionSource, GraphPartitioner graphPartitioner )
    064    {
    065    this( null, phase, partitionSource, graphPartitioner );
    066    }
    067
    068  public RulePartitioner( LogLevel logLevel, PlanPhase phase, PartitionSource partitionSource, GraphPartitioner graphPartitioner )
    069    {
    070    this.logLevel = logLevel;
    071    this.phase = phase;
    072    this.partitionSource = partitionSource;
    073    this.graphPartitioner = graphPartitioner;
    074    }
    075
    076  public RulePartitioner()
    077    {
    078    }
    079
    080  @Override
    081  public PlanPhase getRulePhase()
    082    {
    083    return phase;
    084    }
    085
    086  public PartitionSource getPartitionSource()
    087    {
    088    return partitionSource;
    089    }
    090
    091  protected GraphPartitioner getGraphPartitioner()
    092    {
    093    return graphPartitioner;
    094    }
    095
    096  public Partitions partition( PlannerContext plannerContext, ElementGraph elementGraph )
    097    {
    098    return partition( plannerContext, elementGraph, null );
    099    }
    100
    101  public Partitions partition( PlannerContext plannerContext, ElementGraph elementGraph, Collection<FlowElement> excludes )
    102    {
    103    String[] logLevels = enableLogging( logLevel );
    104
    105    try
    106      {
    107      return performPartition( plannerContext, elementGraph, excludes );
    108      }
    109    finally
    110      {
    111      restoreLogging( logLevels );
    112      }
    113    }
    114
    115  public Partitions performPartition( PlannerContext plannerContext, ElementGraph elementGraph, Collection<FlowElement> excludes )
    116    {
    117    Partitions partitions = getGraphPartitioner().partition( plannerContext, elementGraph, excludes );
    118
    119    if( partitions != null )
    120      partitions.setRulePartitioner( this );
    121
    122    return partitions;
    123    }
    124
    125  @Override
    126  public String getRuleName()
    127    {
    128    return getClass().getSimpleName().replaceAll( "^(.*)[]A-Z][a-z]*Rule$", "$1" );
    129    }
    130
    131  @Override
    132  public String toString()
    133    {
    134    return getRuleName();
    135    }
    136  }