import com.sun.istack.internal.Nullable;

/**
 * <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();
}

import com.google.common.collect.BiMap;
import com.google.common.collect.HashBiMap;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

/**
 * DataType for computing lineage for each column in Wrangler
 * Instance of this type to be sent to platform through API
 */
public class WranglerFieldLevelLineage implements FieldLevelLineage {
  private static final String PARSE_KEY = "all columns";

  private enum Type {
    ADD, DROP, MODIFY, READ, RENAME
  }

  public class WranglerBranchingStepNode implements BranchingTransformStepNode {
    final int stepNumber;
    boolean continueUp;
    boolean continueDown;
    Map<String, Integer> upBranches;
    Map<String, Integer> downBranches;

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

    void putRead(String columnName) {
      upBranches.put(columnName, lineage.get(columnName).size() - 1);
    }

    void putLineage(String columnName) {
      downBranches.put(columnName, lineage.get(columnName).size());
    }

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

    @Override
    public boolean continueBackward() {
      return this.continueDown;
    }

    @Override
    public boolean continueForward() {
      return this.continueUp;
    }

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

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

    // toString()
  }

  private int currentStepNumber;
  private final Set<String> currentColumns;
  private final TransformStep[] steps;
  private final Map<String, List<BranchingTransformStepNode>> lineage;

  public WranglerFieldLevelLineage(int numberOfSteps, List<String> finalColumns) {
    this.currentColumns = new HashSet<>(finalColumns.size());
    this.steps = new WranglerTransformStep[numberOfSteps];
    this.lineage = new HashMap<>(finalColumns.size());
    for (String key : finalColumns) {
      this.currentColumns.add(key);
      this.lineage.put(key, new ArrayList<BranchingTransformStepNode>());
    }
  }

  // store() helpers

  public void store(WranglerStep[] parseTree) {...} // builds lineage data structure

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

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

  // toString()
}

/**
 * <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 {
  List<String> endColumns = new ArrayList<>(oSchema.getFields().size());
  for (Schema.Field field : oSchema.getFields()) {
    endColumns.add(field.getName());
  }
  FieldLevelLineage f = new WranglerFieldLevelLineage(StringUtils.countMatches(config.directives, "\n"),
                                                      endColumns);
  // ((WranglerFieldLevelLineage) f).store(this.parseTree);
  context.recordLineage(f);
}
...