/**
 * <p>
 *   A TransformStep represents a modification done by a Transform.
 * </p>
 */
public interface TransformStep {

  /**
   * @return the name of this modification
   */
  String getName();

  /**
   * @return additional information about this modification
   */
  @Nullable
  String getInformation();
}

import java.util.List;
import java.util.Map;

/**
 * <p>FieldLevelLineage is a DataType for computing lineage for each field in a dataset.
 * An instance of this type can be sent to be sent to platform through API for transform stages.</p>
 */
public interface FieldLevelLineage {

  /**
   * <p>A BranchingTransformStepNode represents a linking between a {@link TransformStep} and a field of data.
   * It contains data about how this TransformStep affected this field.</p>
   */
  interface BranchingTransformStepNode {

    /**
     * @return the index of the TransformStep in {@link #getSteps()}
     */
    int getTransformStepNumber();

    /**
     * @return true if this TransformStep does not add this field
     */
    boolean continueBackward(); // continue down

    /**
     * @return true if this TransformStep does not drop this field
     */
    boolean continueForward();  // continue up

    /**
     * This map should contain every other field that was impacted by this field in this TransformStep
     * @return A map from field name to the index of the next TransformStep using {@link #getLineage()}
     * Usage: getLineage().get(field).get(index).
     */
    Map<String, Integer> getImpactedBranches();  // getUpBranches

    /**
     * This map should contain every other field that impacted this field with this TransformStep
     * @return A map from field name to the index of the previous transformation step in {@link #getLineage()}
     * Usage: getLineage().get(field).get(index).
     */
    Map<String, Integer> getImpactingBranches();  // getDownBranches
  }

  /**
   * @return a list of all TransformSteps executed by this Transform in order.
   */
  List<TransformStep> getSteps();

  /**
   * @return a mapping of field names to a list of {@link BranchingTransformStepNode} in reverse order.
   */
  Map<String, List<BranchingTransformStepNode>> getLineage();
}

/**
 * WranglerFieldLevelLineage is a data type for computing lineage for all columns in Wrangler.
 * An instance of this type to be sent to platform through prepareRun API.
 */
public final class WranglerFieldLevelLineage implements FieldLevelLineage {
  /**
   * A WranglerBranchingStepNode is a data type for linking columns to directives as well as other columns.
   */
  public final class WranglerBranchingStepNode implements BranchingTransformStepNode {
    private final int stepNumber;
    private boolean continueUp;
    private boolean continueDown;
    private final Map<String, Integer> upBranches;
    private final Map<String, Integer> downBranches;

    private WranglerBranchingStepNode(int stepNumber, boolean continueUp, boolean continueDown) {
      this.stepNumber = stepNumber;
      this.continueUp = continueUp;
      this.continueDown = continueDown;
      this.upBranches = new HashMap<>();
      this.downBranches = new HashMap<>();
    }
	
	// Internal functions

    @Override
    public int getTransformStepNumber() {
      return stepNumber;
    }

    @Override
    public boolean continueBackward() {
      return continueUp;
    }

    @Override
    public boolean continueForward() {
      return continueDown;
    }

    @Override
    public Map<String, Integer> getImpactedBranches() {
      return Collections.unmodifiableMap(downBranches);
    }

    @Override
    public Map<String, Integer> getImpactingBranches() {
      return Collections.unmodifiableMap(upBranches);
    }

    @Override
    public String toString() {
      return "(StepNumber: " + stepNumber + ", Continue Backward: " + continueUp + ", " + upBranches +
          ", Continue Forward: " + continueDown + ", " + downBranches + ")";
    }
  }

  private final List<TransformStep> steps;
  private final Map<String, List<BranchingTransformStepNode>> lineage;


  // Internal functions

  // Shows lineage in tree format
  void prettyPrint(String column, boolean forward) {...}

  @Override
  public List<TransformStep> getSteps() {
    return Collections.unmodifiableList(this.steps);
  }

  @Override
  public Map<String, List<BranchingTransformStepNode>> getLineage() {
    return Collections.unmodifiableMap(this.lineage);
  }

  @Override
  public String toString() {
    return "Column Directives: " + lineage + "\n" + "Steps: " + steps;
  }

  /**
   * Builder class for {@link WranglerFieldLevelLineage}.
   */
  public static class Builder {
    // Internal functions
	public Builder(List<String> startColumns, List<String> finalColumns) {...}
    public WranglerFieldLevelLineage build(List<MutationDefinition> parseTree) throws LineageGenerationException {...}
  }
}

/**
 * <p>A FieldLevelLineageStorageNode represents a node of field-level lineage information.</p>
 */
public interface FieldLevelLineageStorageNode {

  /**
   * @return the ID of the pipeline
   */
  ProgramRunId getPipeline();

  /**
   * @return the name of the stage
   */
  String getStage();

  /**
   * @return the name of the field
   */
  String getField();
}

import com.google.common.collect.ArrayListMultimap;
import com.google.common.collect.HashBasedTable;
import com.google.common.collect.ListMultimap;
import com.google.common.collect.Table;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

public final class FieldLevelLineageStorageGraph {
  private final ProgramRunId pipelineId;
  private final PipelinePhase pipeline; // Or another type of graph of the pipeline
  private final Map<String, FieldLevelLineage> stages;
  private final Map<String, String> stageToDataSet;
  private final Map<FieldLevelLineageStorageNode, FieldLevelLineageStorageNode> nodeRetriever;
  private final Table<String, String, DataSetFieldNode> history; // Storage
  private final Map<FieldStepNode, TransformStep> stepInformation; // Storage
  private final ListMultimap<FieldLevelLineageStorageNode, FieldLevelLineageStorageNode> pastEdges; // Storage
  private final ListMultimap<FieldLevelLineageStorageNode, FieldLevelLineageStorageNode> futureEdges; // Storage

  public FieldLevelLineageStorageGraph(ProgramRunId PipelineId, PipelinePhase pipeline,
                                       Map<String, FieldLevelLineage> stages, Map<String, String> stageToDataSet) {
    this.pipelineId = PipelineId;
    this.pipeline = pipeline;
    this.stages = stages;
    this.stageToDataSet = stageToDataSet;
    this.nodeRetriever = new HashMap<>();
    this.history = HashBasedTable.create();
    this.stepInformation = new HashMap<>();
    this.pastEdges = ArrayListMultimap.create();
    this.futureEdges = ArrayListMultimap.create();
  }

  // helpers for make()

  private void make() {...} // makes the graph
}

...
@Override
public void prepareRun(StageSubmitter<TransformContext> context) throws Exception {
  if (context.getContext().getArguments().has("enable.wrangler.field.lineage")) {
    List<String> startColumns, endColumns;
    Schema inputSchema, outputSchema;
    inputSchema = context.getContext().getInputSchema();
    outputSchema = context.getContext().getOutputSchema();
    if (inputSchema == null || inputSchema.getFields() == null) {
      startColumns = new ArrayList<>(0);
    } else {
      startColumns = new ArrayList<>(inputSchema.getFields().size());
      for (Schema.Field field : inputSchema.getFields()) {
        startColumns.add(field.getName());
      }
    }
    if (outputSchema == null || outputSchema.getFields() == null) {
      endColumns = new ArrayList<>(0);
    } else {
      endColumns = new ArrayList<>(outputSchema.getFields().size());
      for (Schema.Field field : outputSchema.getFields()) {
        endColumns.add(field.getName());
      }
    }
    FieldLevelLineage f;
    store = new DefaultTransientStore();
    registry = new CompositeDirectiveRegistry(
      new SystemDirectiveRegistry(),
      new UserDirectiveRegistry(context.getContext())
    );
    RecipeParser directives = new GrammarBasedParser(new MigrateToV2(config.directives).migrate(), registry);
    ExecutorContext ctx = new WranglerPipelineContext(ExecutorContext.Environment.TRANSFORM,
                                                      context.getContext(), store);
    try {
      // Create the pipeline executor with context being set.
      pipeline = new RecipePipelineExecutor();
      pipeline.initialize(directives, ctx);
    } catch (Exception e) {
      throw new Exception(
        String.format("Stage:%s - %s", context.getContext().getStageName(), e.getMessage())
      );
    }
    try {
      f = (new WranglerFieldLevelLineage.Builder(startColumns, endColumns)).build(pipeline.lineage());
    } catch (LineageGenerationException e) {
      LOG.warn(
        String.format("Unable to generate lineage: %s", e.getMessage())
      );
      f = null;
    }
    LOG.info(
      "Field-level lineage generated"
    );
    // context.recordLineage(f);
  }
}
...