ColumnLineageGraph.java

Go to the documentation of this file.

// Copyright 2012 Cloudera Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package com.cloudera.impala.analysis;

import java.text.SimpleDateFormat;
import java.util.Collection;
import java.util.Date;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.TreeSet;

import org.json.simple.JSONArray;
import org.json.simple.JSONObject;
import org.json.simple.JSONValue;
import org.json.simple.parser.JSONParser;
import org.json.simple.parser.ParseException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.cloudera.impala.catalog.Table;
import com.cloudera.impala.common.Id;
import com.cloudera.impala.common.IdGenerator;
import com.cloudera.impala.thrift.TQueryCtx;
import com.google.common.base.Joiner;
import com.google.common.base.Preconditions;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Sets;
import com.google.common.hash.Hasher;
import com.google.common.hash.Hashing;

final class Vertex implements Comparable<Vertex> {
  // Unique identifier of this vertex.
  private final VertexId id_;

  private final String type_ = "COLUMN";

  // A fully-qualified column name or the label of a result expr
  private final String label_;

  public Vertex(VertexId id, String label) {
    Preconditions.checkNotNull(id);
    Preconditions.checkNotNull(label);
    id_ = id;
    label_ = label;
  }
  public VertexId getVertexId() { return id_; }
  public String getLabel() { return label_; }
  public String getType() { return type_; }

  @Override
  public String toString() { return "(" + id_ + ":" + type_ + ":" + label_ + ")"; }

  public Map toJson() {
    // Use a LinkedHashMap to generate a strict ordering of elements.
    Map obj = new LinkedHashMap();
    obj.put("id", id_.asInt());
    obj.put("vertexType", type_);
    obj.put("vertexId", label_);
    return obj;
  }

  public static Vertex fromJsonObj(JSONObject obj) {
    int id = ((Long) obj.get("id")).intValue();
    String label = (String) obj.get("vertexId");
    return new Vertex(new VertexId(id), label);
  }

  @Override
  public boolean equals(Object obj) {
    if (obj == null) return false;
    if (obj.getClass() != this.getClass()) return false;
    Vertex vertex = (Vertex) obj;
    return this.id_.equals(vertex.id_) &&
        this.label_.equals(vertex.label_);
  }

  public int compareTo(Vertex cmp) { return this.id_.compareTo(cmp.id_); }

  @Override
  public int hashCode() { return id_.hashCode(); }
}

class VertexId extends Id<VertexId> {
  protected VertexId(int id) {
    super(id);
  }
  public static IdGenerator<VertexId> createGenerator() {
    return new IdGenerator<VertexId>() {
      @Override
      public VertexId getNextId() { return new VertexId(nextId_++); }
      @Override
      public VertexId getMaxId() { return new VertexId(nextId_ - 1); }
    };
  }
}

final class MultiEdge {
  public static enum EdgeType {
    PROJECTION, PREDICATE
  }
  private final Set<Vertex> sources_;
  private final Set<Vertex> targets_;
  private final EdgeType edgeType_;

  public MultiEdge(Set<Vertex> sources, Set<Vertex> targets, EdgeType type) {
    sources_ = sources;
    targets_ = targets;
    edgeType_ = type;
  }

  @Override
  public String toString() {
    StringBuilder builder = new StringBuilder();
    Joiner joiner = Joiner.on(",");
    builder.append("Sources: [");
    builder.append(joiner.join(sources_) + "]\n");
    builder.append("Targets: [");
    builder.append(joiner.join(targets_) + "]\n");
    builder.append("Type: " + edgeType_);
    return builder.toString();
  }

  public Map toJson() {
    Map obj = new LinkedHashMap();
    // Add sources
    JSONArray sourceIds = new JSONArray();
    for (Vertex vertex: sources_) {
      sourceIds.add(vertex.getVertexId());
    }
    obj.put("sources", sourceIds);
    // Add targets
    JSONArray targetIds = new JSONArray();
    for (Vertex vertex: targets_) {
      targetIds.add(vertex.getVertexId());
    }
    obj.put("targets", targetIds);
    obj.put("edgeType", edgeType_.toString());
    return obj;
  }

  @Override
  public boolean equals(Object obj) {
    if (obj == null) return false;
    if (obj.getClass() != this.getClass()) return false;
    MultiEdge edge = (MultiEdge) obj;
    return edge.sources_.equals(this.sources_) &&
        edge.targets_.equals(this.targets_) &&
        edge.edgeType_ == this.edgeType_;
  }
}

public class ColumnLineageGraph {
  private final static Logger LOG = LoggerFactory.getLogger(ColumnLineageGraph.class);
  // Query statement
  private String queryStr_;

  // Name of the user that issued this query
  private String user_;

  private final List<Expr> resultDependencyPredicates_ = Lists.newArrayList();

  private final List<MultiEdge> edges_ = Lists.newArrayList();

  // Timestamp in seconds since epoch (GMT) this query was submitted for execution.
  private long timestamp_;

  // Map of Vertex labels to Vertex objects.
  private final Map<String, Vertex> vertices_ = Maps.newHashMap();

  // Map of Vertex ids to Vertex objects. Used primarily during the construction of the
  // ColumnLineageGraph from a serialized JSON object.
  private final Map<VertexId, Vertex> idToVertexMap_ = Maps.newHashMap();

  // For an INSERT or a CTAS, these are the columns of the
  // destination table plus any partitioning columns (when dynamic partitioning is used).
  // For a SELECT stmt, they are the labels of the result exprs.
  private final List<String> targetColumnLabels_ = Lists.newArrayList();

  // Repository for tuple and slot descriptors for this query. Use it to construct the
  // column lineage graph.
  private DescriptorTable descTbl_;

  private final IdGenerator<VertexId> vertexIdGenerator = VertexId.createGenerator();

  public ColumnLineageGraph() { }

  private ColumnLineageGraph(String stmt, String user, long timestamp) {
    queryStr_ = stmt;
    user_ = user;
    timestamp_ = timestamp;
  }

  private void setVertices(Set<Vertex> vertices) {
    for (Vertex vertex: vertices) {
      vertices_.put(vertex.getLabel(), vertex);
      idToVertexMap_.put(vertex.getVertexId(), vertex);
    }
  }

  private MultiEdge createMultiEdge(Set<String> targets, Set<String> sources,
      MultiEdge.EdgeType type) {
    Set<Vertex> targetVertices = Sets.newHashSet();
    for (String target: targets) {
      targetVertices.add(createVertex(target));
    }
    Set<Vertex> sourceVertices = Sets.newHashSet();
    for (String source: sources) {
      sourceVertices.add(createVertex(source));
    }
    MultiEdge edge = new MultiEdge(sourceVertices, targetVertices, type);
    edges_.add(edge);
    return edge;
  }

  private Vertex createVertex(String label) {
    Vertex newVertex = vertices_.get(label);
    if (newVertex != null) return newVertex;
    newVertex = new Vertex(vertexIdGenerator.getNextId(), label);
    vertices_.put(newVertex.getLabel(), newVertex);
    idToVertexMap_.put(newVertex.getVertexId(), newVertex);
    return newVertex;
  }

  public void computeLineageGraph(List<Expr> resultExprs, Analyzer rootAnalyzer) {
    init(rootAnalyzer);
    computeProjectionDependencies(resultExprs);
    computeResultPredicateDependencies(rootAnalyzer);
  }

  private void init(Analyzer analyzer) {
    Preconditions.checkNotNull(analyzer);
    Preconditions.checkState(analyzer.isRootAnalyzer());
    TQueryCtx queryCtx = analyzer.getQueryCtx();
    if (queryCtx.request.isSetRedacted_stmt()) {
      queryStr_ = queryCtx.request.redacted_stmt;
    } else {
      queryStr_ = queryCtx.request.stmt;
    }
    Preconditions.checkNotNull(queryStr_);
    SimpleDateFormat df = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss");
    try {
      timestamp_ = df.parse(queryCtx.now_string).getTime() / 1000;
    } catch (java.text.ParseException e) {
      LOG.error("Error parsing timestamp value: " + queryCtx.now_string +
          " " + e.getMessage());
      timestamp_ = new Date().getTime() / 1000;
    }
    descTbl_ = analyzer.getDescTbl();
    user_ = analyzer.getUser().getName();
  }

  private void computeProjectionDependencies(List<Expr> resultExprs) {
    Preconditions.checkNotNull(resultExprs);
    Preconditions.checkState(!resultExprs.isEmpty());
    Preconditions.checkState(resultExprs.size() == targetColumnLabels_.size());
    for (int i = 0; i < resultExprs.size(); ++i) {
      Expr expr = resultExprs.get(i);
      Set<String> sourceBaseCols = Sets.newHashSet();
      List<Expr> dependentExprs = Lists.newArrayList();
      getSourceBaseCols(expr, sourceBaseCols, dependentExprs, false);
      Set<String> targets = Sets.newHashSet(targetColumnLabels_.get(i));
      createMultiEdge(targets, sourceBaseCols, MultiEdge.EdgeType.PROJECTION);
      if (!dependentExprs.isEmpty()) {
        // We have additional exprs that 'expr' has a predicate dependency on.
        // Gather the transitive predicate dependencies of 'expr' based on its direct
        // predicate dependencies. For each direct predicate dependency p, 'expr' is
        // transitively predicate dependent on all exprs that p is projection and
        // predicate dependent on.
        Set<String> predicateBaseCols = Sets.newHashSet();
        for (Expr dependentExpr: dependentExprs) {
          getSourceBaseCols(dependentExpr, predicateBaseCols, null, true);
        }
        createMultiEdge(targets, predicateBaseCols, MultiEdge.EdgeType.PREDICATE);
      }
    }
  }

  private void computeResultPredicateDependencies(Analyzer analyzer) {
    List<Expr> conjuncts = analyzer.getConjuncts();
    for (Expr expr: conjuncts) {
      if (expr.isAuxExpr()) continue;
      resultDependencyPredicates_.add(expr);
    }
    Set<String> predicateBaseCols = Sets.newHashSet();
    for (Expr expr: resultDependencyPredicates_) {
      getSourceBaseCols(expr, predicateBaseCols, null, true);
    }
    if (predicateBaseCols.isEmpty()) return;
    Set<String> targets = Sets.newHashSet(targetColumnLabels_);
    createMultiEdge(targets, predicateBaseCols, MultiEdge.EdgeType.PREDICATE);
  }

  private void getSourceBaseCols(Expr expr, Set<String> sourceBaseCols,
      List<Expr> directPredDeps, boolean traversePredDeps) {
    List<Expr> exprsToTraverse = getProjectionDeps(expr);
    List<Expr> predicateDepExprs = getPredicateDeps(expr);
    if (directPredDeps != null) directPredDeps.addAll(predicateDepExprs);
    if (traversePredDeps) exprsToTraverse.addAll(predicateDepExprs);
    List<SlotId> slotIds = Lists.newArrayList();
    for (Expr e: exprsToTraverse) {
      e.getIds(null, slotIds);
    }
    for (SlotId slotId: slotIds) {
      SlotDescriptor slotDesc = descTbl_.getSlotDesc(slotId);
      List<Expr> sourceExprs = slotDesc.getSourceExprs();
      // TODO: Add lineage support for nested types.
      if (sourceExprs.isEmpty() && slotDesc.getColumn() != null) {
        // slot should correspond to a materialized base table column
        Preconditions.checkState(slotDesc.getParent().isMaterialized()
            && slotDesc.getParent().getTable() != null
            && slotDesc.getColumn() != null);
        String colName = slotDesc.getParent().getTableName() + "." +
            slotDesc.getColumn().getName();
        sourceBaseCols.add(colName);
      } else {
        for (Expr sourceExpr: sourceExprs) {
          getSourceBaseCols(sourceExpr, sourceBaseCols, directPredDeps,
              traversePredDeps);
        }
      }
    }
  }

  private List<Expr> getProjectionDeps(Expr e) {
    Preconditions.checkNotNull(e);
    List<Expr> outputExprs = Lists.newArrayList();
    if (e instanceof AnalyticExpr) {
      AnalyticExpr analytic = (AnalyticExpr) e;
      outputExprs.addAll(analytic.getChildren().subList(0,
          analytic.getFnCall().getParams().size()));
    } else {
      outputExprs.add(e);
    }
    return outputExprs;
  }

  private List<Expr> getPredicateDeps(Expr e) {
    Preconditions.checkNotNull(e);
    List<Expr> outputExprs = Lists.newArrayList();
    if (e instanceof AnalyticExpr) {
      AnalyticExpr analyticExpr = (AnalyticExpr) e;
      outputExprs.addAll(analyticExpr.getPartitionExprs());
      for (OrderByElement orderByElem: analyticExpr.getOrderByElements()) {
        outputExprs.add(orderByElem.getExpr());
      }
    }
    return outputExprs;
  }

  public void addDependencyPredicates(Collection<Expr> exprs) {
    resultDependencyPredicates_.addAll(exprs);
  }

  public String toJson() {
    if (queryStr_ == null) return "";
    Map obj = new LinkedHashMap();
    obj.put("queryText", queryStr_);
    obj.put("hash", getQueryHash(queryStr_));
    obj.put("user", user_);
    obj.put("timestamp", timestamp_);
    // Add edges
    JSONArray edges = new JSONArray();
    for (MultiEdge edge: edges_) {
      edges.add(edge.toJson());
    }
    obj.put("edges", edges);
    // Add vertices
    TreeSet<Vertex> sortedVertices = Sets.newTreeSet(vertices_.values());
    JSONArray vertices = new JSONArray();
    for (Vertex vertex: sortedVertices) {
      vertices.add(vertex.toJson());
    }
    obj.put("vertices", vertices);
    return JSONValue.toJSONString(obj);
  }

  private String getQueryHash(String queryStr) {
    Hasher hasher = Hashing.md5().newHasher();
    hasher.putString(queryStr);
    return hasher.hash().toString();
  }

  public static ColumnLineageGraph createFromJSON(String json) {
    if (json == null || json.isEmpty()) return null;
    JSONParser parser = new JSONParser();
    Object obj = null;
    try {
      obj = parser.parse(json);
    } catch (ParseException e) {
      LOG.error("Error parsing serialized column lineage graph: " + e.getMessage());
      return null;
    }
    if (!(obj instanceof JSONObject)) return null;
    JSONObject jsonObj = (JSONObject) obj;
    String stmt = (String) jsonObj.get("queryText");
    String hash = (String) jsonObj.get("hash");
    String user = (String) jsonObj.get("user");
    long timestamp = (Long) jsonObj.get("timestamp");
    ColumnLineageGraph graph = new ColumnLineageGraph(stmt, user, timestamp);
    JSONArray serializedVertices = (JSONArray) jsonObj.get("vertices");
    Set<Vertex> vertices = Sets.newHashSet();
    for (int i = 0; i < serializedVertices.size(); ++i) {
      Vertex v = Vertex.fromJsonObj((JSONObject) serializedVertices.get(i));
      vertices.add(v);
    }
    graph.setVertices(vertices);
    JSONArray serializedEdges = (JSONArray) jsonObj.get("edges");
    for (int i = 0; i < serializedEdges.size(); ++i) {
      MultiEdge e =
          graph.createMultiEdgeFromJSONObj((JSONObject) serializedEdges.get(i));
      graph.edges_.add(e);
    }
    return graph;
  }

  private MultiEdge createMultiEdgeFromJSONObj(JSONObject jsonEdge) {
    Preconditions.checkNotNull(jsonEdge);
    JSONArray sources = (JSONArray) jsonEdge.get("sources");
    Set<Vertex> sourceVertices = getVerticesFromJSONArray(sources);
    JSONArray targets = (JSONArray) jsonEdge.get("targets");
    Set<Vertex> targetVertices = getVerticesFromJSONArray(targets);
    MultiEdge.EdgeType type =
        MultiEdge.EdgeType.valueOf((String) jsonEdge.get("edgeType"));
    return new MultiEdge(sourceVertices, targetVertices, type);
  }

  private Set<Vertex> getVerticesFromJSONArray(JSONArray vertexIdArray) {
    Set<Vertex> vertices = Sets.newHashSet();
    for (int i = 0; i < vertexIdArray.size(); ++i) {
      int sourceId = ((Long) vertexIdArray.get(i)).intValue();
      Vertex sourceVertex = idToVertexMap_.get(new VertexId(sourceId));
      Preconditions.checkNotNull(sourceVertex);
      vertices.add(sourceVertex);
    }
    return vertices;
  }

  @Override
  public boolean equals(Object obj) {
    if (obj == null) return false;
    if (obj.getClass() != this.getClass()) return false;
    ColumnLineageGraph g = (ColumnLineageGraph) obj;
    if (!this.vertices_.equals(g.vertices_) ||
        !this.edges_.equals(g.edges_)) {
      return false;
    }
    return true;
  }

  public String debugString() {
    StringBuilder builder = new StringBuilder();
    for (MultiEdge edge: edges_) {
      builder.append(edge.toString() + "\n");
    }
    builder.append(toJson());
    return builder.toString();
  }

  public void addTargetColumnLabels(Collection<String> columnLabels) {
    Preconditions.checkNotNull(columnLabels);
    targetColumnLabels_.addAll(columnLabels);
  }

  public void addTargetColumnLabels(Table dstTable) {
    Preconditions.checkNotNull(dstTable);
    String tblFullName = dstTable.getFullName();
    for (String columnName: dstTable.getColumnNames()) {
      targetColumnLabels_.add(tblFullName + "." + columnName);
    }
  }
}