/**
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you 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 org.apache.hadoop.hive.ql.optimizer.lineage;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.Map;
import java.util.Set;
import org.apache.commons.lang.StringUtils;
import org.apache.hadoop.hive.common.ObjectPair;
import org.apache.hadoop.hive.ql.exec.ColumnInfo;
import org.apache.hadoop.hive.ql.exec.FileSinkOperator;
import org.apache.hadoop.hive.ql.exec.Operator;
import org.apache.hadoop.hive.ql.exec.SelectOperator;
import org.apache.hadoop.hive.ql.hooks.LineageInfo;
import org.apache.hadoop.hive.ql.hooks.LineageInfo.BaseColumnInfo;
import org.apache.hadoop.hive.ql.hooks.LineageInfo.Dependency;
import org.apache.hadoop.hive.ql.hooks.LineageInfo.Predicate;
import org.apache.hadoop.hive.ql.lib.NodeProcessorCtx;
import org.apache.hadoop.hive.ql.metadata.Table;
import org.apache.hadoop.hive.ql.parse.ParseContext;
import org.apache.hadoop.hive.ql.plan.OperatorDesc;
/**
* This class contains the lineage context that is passed
* while walking the operator tree in Lineage. The context
* contains the LineageInfo structure that is passed to the
* pre-execution hooks.
*/
public class LineageCtx implements NodeProcessorCtx {
public static class Index {
/**
* The map contains an index from the (operator, columnInfo) to the
* dependency vector for that tuple. This is used to generate the
* dependency vectors during the walk of the operator tree.
*/
private final Map<Operator<? extends OperatorDesc>,
LinkedHashMap<ColumnInfo, Dependency>> depMap;
/**
* A map from operator to the conditions strings.
*/
private final Map<Operator<? extends OperatorDesc>, Set<Predicate>> condMap;
/**
* A map from a final select operator id to the select operator
* and the corresponding target table in case an insert into query.
*/
private LinkedHashMap<String, ObjectPair<SelectOperator, Table>> finalSelectOps;
/**
* Constructor.
*/
public Index() {
depMap =
new LinkedHashMap<Operator<? extends OperatorDesc>,
LinkedHashMap<ColumnInfo, Dependency>>();
condMap = new HashMap<Operator<? extends OperatorDesc>, Set<Predicate>>();
finalSelectOps =
new LinkedHashMap<String, ObjectPair<SelectOperator, Table>>();
}
/**
* Gets the dependency for an operator, columninfo tuple.
* @param op The operator whose dependency is being inspected.
* @param col The column info whose dependency is being inspected.
* @return Dependency for that particular operator, columninfo tuple.
* null if no dependency is found.
*/
public Dependency getDependency(Operator<? extends OperatorDesc> op,
ColumnInfo col) {
Map<ColumnInfo, Dependency> colMap = depMap.get(op);
if (colMap == null) {
return null;
}
return colMap.get(col);
}
/**
* Gets the dependency for a tuple of an operator,
* and a ColumnInfo with specified internal name.
*
* @param op The operator whose dependency is being inspected.
* @param internalName The internal name of the column info
* @return Dependency for that particular operator, ColumnInfo tuple.
* null if no dependency is found.
*/
public Dependency getDependency(
Operator<? extends OperatorDesc> op, String internalName) {
Map<ColumnInfo, Dependency> colMap = depMap.get(op);
if (colMap != null) {
for (Map.Entry<ColumnInfo, Dependency> e: colMap.entrySet()) {
if (e.getKey().getInternalName().equals(internalName)) {
return e.getValue();
}
}
}
return null;
}
/**
* Puts the dependency for an operator, columninfo tuple.
* @param op The operator whose dependency is being inserted.
* @param col The column info whose dependency is being inserted.
* @param dep The dependency.
*/
public void putDependency(Operator<? extends OperatorDesc> op,
ColumnInfo col, Dependency dep) {
LinkedHashMap<ColumnInfo, Dependency> colMap = depMap.get(op);
if (colMap == null) {
colMap = new LinkedHashMap<ColumnInfo, Dependency>();
depMap.put(op, colMap);
}
colMap.put(col, dep);
}
/**
* Merges the new dependencies in dep to the existing dependencies
* of (op, ci).
*
* @param op The operator of the column whose dependency is being modified.
* @param ci The column info of the associated column.
* @param dep The new dependency.
*/
public void mergeDependency(Operator<? extends OperatorDesc> op,
ColumnInfo ci, Dependency dep) {
Dependency old_dep = getDependency(op, ci);
if (old_dep == null) {
putDependency(op, ci, dep);
} else {
LineageInfo.DependencyType new_type =
LineageCtx.getNewDependencyType(old_dep.getType(),
LineageInfo.DependencyType.EXPRESSION);
old_dep.setType(new_type);
Set<BaseColumnInfo> bci_set = new LinkedHashSet<BaseColumnInfo>(old_dep.getBaseCols());
bci_set.addAll(dep.getBaseCols());
old_dep.setBaseCols(bci_set);
// TODO: Fix the expressions later.
old_dep.setExpr(null);
}
}
public Map<ColumnInfo, Dependency> getDependencies(Operator<? extends OperatorDesc> op) {
return depMap.get(op);
}
public void addPredicate(Operator<? extends OperatorDesc> op, Predicate cond) {
Set<Predicate> conds = condMap.get(op);
if (conds == null) {
conds = new LinkedHashSet<Predicate>();
condMap.put(op, conds);
}
for (Predicate p: conds) {
if (StringUtils.equals(cond.getExpr(), p.getExpr())) {
p.getBaseCols().addAll(cond.getBaseCols());
return;
}
}
conds.add(cond);
}
public void copyPredicates(Operator<? extends OperatorDesc> srcOp,
Operator<? extends OperatorDesc> tgtOp) {
Set<Predicate> conds = getPredicates(srcOp);
if (conds != null) {
for (Predicate cond: conds) {
addPredicate(tgtOp, cond);
}
}
}
public Set<Predicate> getPredicates(Operator<? extends OperatorDesc> op) {
return condMap.get(op);
}
public void addFinalSelectOp(
SelectOperator sop, Operator<? extends OperatorDesc> sinkOp) {
String operatorId = sop.getOperatorId();
if (!finalSelectOps.containsKey(operatorId)) {
Table table = null;
if (sinkOp instanceof FileSinkOperator) {
FileSinkOperator fso = (FileSinkOperator) sinkOp;
table = fso.getConf().getTable();
}
finalSelectOps.put(operatorId,
new ObjectPair<SelectOperator, Table>(sop, table));
}
}
public LinkedHashMap<String,
ObjectPair<SelectOperator, Table>> getFinalSelectOps() {
return finalSelectOps;
}
public void clear() {
finalSelectOps.clear();
depMap.clear();
condMap.clear();
}
}
/**
* The map contains an index from the (operator, columnInfo) to the
* dependency vector for that tuple. This is used to generate the
* dependency vectors during the walk of the operator tree.
*/
private final Index index;
/**
* Parse context to get to the table metadata information.
*/
private final ParseContext pctx;
/**
* Constructor.
*
* @param pctx The parse context that is used to get table metadata information.
* @param index The dependency map.
*/
public LineageCtx(ParseContext pctx, Index index) {
this.index = index;
this.pctx = pctx;
}
/**
* Gets the parse context.
*
* @return ParseContext
*/
public ParseContext getParseCtx() {
return pctx;
}
/**
* Gets the dependency index.
*
* @return Index
*/
public Index getIndex() {
return index;
}
/**
* Gets the new dependency type by comparing the old dependency type and the
* current dependency type. The current dependency type is the dependency imposed
* by the current expression. Typically the dependency type is computed using
* the following rules:
* SCRIPT - In case anywhere in the lineage tree there was a script operator, otherwise
* EXPRESSION - In case anywhere in the lineage tree a union,
* udf, udaf or udtf was done, otherwise
* SIMPLE - This captures direct column copies.
*
* @param old_type The old dependency type.
* @param curr_type The current operators dependency type.
* @return the dependency type
*/
public static LineageInfo.DependencyType getNewDependencyType(
LineageInfo.DependencyType old_type, LineageInfo.DependencyType curr_type) {
if (old_type == LineageInfo.DependencyType.SCRIPT ||
curr_type == LineageInfo.DependencyType.SCRIPT) {
return LineageInfo.DependencyType.SCRIPT;
}
if (old_type == LineageInfo.DependencyType.EXPRESSION ||
curr_type == LineageInfo.DependencyType.EXPRESSION) {
return LineageInfo.DependencyType.EXPRESSION;
}
return LineageInfo.DependencyType.SIMPLE;
}
}