private void handleSampling(DriverContext context, MapWork mWork, JobConf job, HiveConf conf)
throws Exception {
assert mWork.getAliasToWork().keySet().size() == 1;
String alias = mWork.getAliases().get(0);
Operator<?> topOp = mWork.getAliasToWork().get(alias);
PartitionDesc partDesc = mWork.getAliasToPartnInfo().get(alias);
ArrayList<String> paths = mWork.getPaths();
ArrayList<PartitionDesc> parts = mWork.getPartitionDescs();
List<Path> inputPaths = new ArrayList<Path>(paths.size());
for (String path : paths) {
inputPaths.add(new Path(path));
}
Path tmpPath = context.getCtx().getExternalTmpPath(inputPaths.get(0));
Path partitionFile = new Path(tmpPath, ".partitions");
ShimLoader.getHadoopShims().setTotalOrderPartitionFile(job, partitionFile);
PartitionKeySampler sampler = new PartitionKeySampler();
if (mWork.getSamplingType() == MapWork.SAMPLING_ON_PREV_MR) {
console.printInfo("Use sampling data created in previous MR");
// merges sampling data from previous MR and make partition keys for total sort
for (Path path : inputPaths) {
FileSystem fs = path.getFileSystem(job);
for (FileStatus status : fs.globStatus(new Path(path, ".sampling*"))) {
sampler.addSampleFile(status.getPath(), job);
}
}
} else if (mWork.getSamplingType() == MapWork.SAMPLING_ON_START) {
console.printInfo("Creating sampling data..");
assert topOp instanceof TableScanOperator;
TableScanOperator ts = (TableScanOperator) topOp;
FetchWork fetchWork;
if (!partDesc.isPartitioned()) {
assert paths.size() == 1;
fetchWork = new FetchWork(inputPaths.get(0), partDesc.getTableDesc());
} else {
fetchWork = new FetchWork(inputPaths, parts, partDesc.getTableDesc());
}
fetchWork.setSource(ts);
// random sampling
FetchOperator fetcher = PartitionKeySampler.createSampler(fetchWork, conf, job, ts);
try {
ts.initialize(conf, new ObjectInspector[] {fetcher.getOutputObjectInspector()});
OperatorUtils.setChildrenCollector(ts.getChildOperators(), sampler);
while (fetcher.pushRow()) {}
} finally {
fetcher.clearFetchContext();
}
} else {
throw new IllegalArgumentException("Invalid sampling type " + mWork.getSamplingType());
}
sampler.writePartitionKeys(partitionFile, conf, job);
}
// Remove RS and SEL introduced by enforce bucketing/sorting config
// Convert PARENT -> RS -> SEL -> FS to PARENT -> FS
private boolean removeRSInsertedByEnforceBucketing(FileSinkOperator fsOp) {
Set<ReduceSinkOperator> reduceSinks =
OperatorUtils.findOperatorsUpstream(fsOp, ReduceSinkOperator.class);
Operator<? extends OperatorDesc> rsToRemove = null;
List<ReduceSinkOperator> rsOps =
parseCtx.getReduceSinkOperatorsAddedByEnforceBucketingSorting();
boolean found = false;
// iterate through all RS and locate the one introduce by enforce bucketing
for (ReduceSinkOperator reduceSink : reduceSinks) {
for (ReduceSinkOperator rsOp : rsOps) {
if (reduceSink.equals(rsOp)) {
rsToRemove = reduceSink;
found = true;
break;
}
}
if (found) {
break;
}
}
// iF RS is found remove it and its child (EX) and connect its parent
// and grand child
if (found) {
Operator<? extends OperatorDesc> rsParent = rsToRemove.getParentOperators().get(0);
Operator<? extends OperatorDesc> rsChild = rsToRemove.getChildOperators().get(0);
Operator<? extends OperatorDesc> rsGrandChild = rsChild.getChildOperators().get(0);
if (rsChild instanceof SelectOperator) {
// if schema size cannot be matched, then it could be because of constant folding
// converting partition column expression to constant expression. The constant
// expression will then get pruned by column pruner since it will not reference to
// any columns.
if (rsParent.getSchema().getSignature().size()
!= rsChild.getSchema().getSignature().size()) {
return false;
}
rsParent.getChildOperators().clear();
rsParent.getChildOperators().add(rsGrandChild);
rsGrandChild.getParentOperators().clear();
rsGrandChild.getParentOperators().add(rsParent);
LOG.info(
"Removed "
+ rsToRemove.getOperatorId()
+ " and "
+ rsChild.getOperatorId()
+ " as it was introduced by enforce bucketing/sorting.");
}
}
return true;
}
@Override
@SuppressWarnings("unchecked")
public void init(JobConf job, OutputCollector output, Reporter reporter) throws Exception {
perfLogger.PerfLogBegin(CLASS_NAME, PerfLogger.SPARK_INIT_OPERATORS);
super.init(job, output, reporter);
rowObjectInspector = new ObjectInspector[Byte.MAX_VALUE];
ObjectInspector[] valueObjectInspector = new ObjectInspector[Byte.MAX_VALUE];
ObjectInspector keyObjectInspector;
ReduceWork gWork = Utilities.getReduceWork(job);
reducer = gWork.getReducer();
vectorized = gWork.getVectorMode();
reducer.setParentOperators(null); // clear out any parents as reducer is the
// root
isTagged = gWork.getNeedsTagging();
try {
keyTableDesc = gWork.getKeyDesc();
inputKeyDeserializer = ReflectionUtils.newInstance(keyTableDesc.getDeserializerClass(), null);
SerDeUtils.initializeSerDe(inputKeyDeserializer, null, keyTableDesc.getProperties(), null);
keyObjectInspector = inputKeyDeserializer.getObjectInspector();
valueTableDesc = new TableDesc[gWork.getTagToValueDesc().size()];
if (vectorized) {
final int maxTags = gWork.getTagToValueDesc().size();
keyStructInspector = (StructObjectInspector) keyObjectInspector;
batches = new VectorizedRowBatch[maxTags];
valueStructInspectors = new StructObjectInspector[maxTags];
valueStringWriters = new List[maxTags];
keysColumnOffset = keyStructInspector.getAllStructFieldRefs().size();
buffer = new DataOutputBuffer();
}
for (int tag = 0; tag < gWork.getTagToValueDesc().size(); tag++) {
// We should initialize the SerDe with the TypeInfo when available.
valueTableDesc[tag] = gWork.getTagToValueDesc().get(tag);
inputValueDeserializer[tag] =
ReflectionUtils.newInstance(valueTableDesc[tag].getDeserializerClass(), null);
SerDeUtils.initializeSerDe(
inputValueDeserializer[tag], null, valueTableDesc[tag].getProperties(), null);
valueObjectInspector[tag] = inputValueDeserializer[tag].getObjectInspector();
ArrayList<ObjectInspector> ois = new ArrayList<ObjectInspector>();
if (vectorized) {
/* vectorization only works with struct object inspectors */
valueStructInspectors[tag] = (StructObjectInspector) valueObjectInspector[tag];
ObjectPair<VectorizedRowBatch, StandardStructObjectInspector> pair =
VectorizedBatchUtil.constructVectorizedRowBatch(
keyStructInspector,
valueStructInspectors[tag],
gWork.getVectorScratchColumnTypeMap());
batches[tag] = pair.getFirst();
final int totalColumns =
keysColumnOffset + valueStructInspectors[tag].getAllStructFieldRefs().size();
valueStringWriters[tag] = new ArrayList<VectorExpressionWriter>(totalColumns);
valueStringWriters[tag].addAll(
Arrays.asList(
VectorExpressionWriterFactory.genVectorStructExpressionWritables(
keyStructInspector)));
valueStringWriters[tag].addAll(
Arrays.asList(
VectorExpressionWriterFactory.genVectorStructExpressionWritables(
valueStructInspectors[tag])));
rowObjectInspector[tag] = pair.getSecond();
} else {
ois.add(keyObjectInspector);
ois.add(valueObjectInspector[tag]);
// reducer.setGroupKeyObjectInspector(keyObjectInspector);
rowObjectInspector[tag] =
ObjectInspectorFactory.getStandardStructObjectInspector(
Utilities.reduceFieldNameList, ois);
}
}
} catch (Exception e) {
throw new RuntimeException(e);
}
ExecMapperContext execContext = new ExecMapperContext(job);
localWork = gWork.getMapRedLocalWork();
execContext.setJc(jc);
execContext.setLocalWork(localWork);
reducer.passExecContext(execContext);
reducer.setReporter(rp);
OperatorUtils.setChildrenCollector(
Arrays.<Operator<? extends OperatorDesc>>asList(reducer), output);
// initialize reduce operator tree
try {
LOG.info(reducer.dump(0));
reducer.initialize(jc, rowObjectInspector);
if (localWork != null) {
for (Operator<? extends OperatorDesc> dummyOp : localWork.getDummyParentOp()) {
dummyOp.setExecContext(execContext);
dummyOp.initialize(jc, null);
}
}
} catch (Throwable e) {
abort = true;
if (e instanceof OutOfMemoryError) {
// Don't create a new object if we are already out of memory
throw (OutOfMemoryError) e;
} else {
throw new RuntimeException("Reduce operator initialization failed", e);
}
}
perfLogger.PerfLogEnd(CLASS_NAME, PerfLogger.SPARK_INIT_OPERATORS);
}
public ReduceSinkOperator getReduceSinkOp(
List<Integer> partitionPositions,
List<Integer> sortPositions,
List<Integer> sortOrder,
List<Integer> sortNullOrder,
ArrayList<ExprNodeDesc> allCols,
ArrayList<ExprNodeDesc> bucketColumns,
int numBuckets,
Operator<? extends OperatorDesc> parent,
AcidUtils.Operation writeType)
throws SemanticException {
// Order of KEY columns
// 1) Partition columns
// 2) Bucket number column
// 3) Sort columns
Set<Integer> keyColsPosInVal = Sets.newLinkedHashSet();
ArrayList<ExprNodeDesc> keyCols = Lists.newArrayList();
List<Integer> newSortOrder = Lists.newArrayList();
List<Integer> newSortNullOrder = Lists.newArrayList();
int numPartAndBuck = partitionPositions.size();
keyColsPosInVal.addAll(partitionPositions);
if (!bucketColumns.isEmpty()
|| writeType == Operation.DELETE
|| writeType == Operation.UPDATE) {
keyColsPosInVal.add(-1);
numPartAndBuck += 1;
}
keyColsPosInVal.addAll(sortPositions);
// by default partition and bucket columns are sorted in ascending order
Integer order = 1;
if (sortOrder != null && !sortOrder.isEmpty()) {
if (sortOrder.get(0).intValue() == 0) {
order = 0;
}
}
for (int i = 0; i < numPartAndBuck; i++) {
newSortOrder.add(order);
}
newSortOrder.addAll(sortOrder);
String orderStr = "";
for (Integer i : newSortOrder) {
if (i.intValue() == 1) {
orderStr += "+";
} else {
orderStr += "-";
}
}
// if partition and bucket columns are sorted in ascending order, by default
// nulls come first; otherwise nulls come last
Integer nullOrder = order == 1 ? 0 : 1;
if (sortNullOrder != null && !sortNullOrder.isEmpty()) {
if (sortNullOrder.get(0).intValue() == 0) {
nullOrder = 0;
} else {
nullOrder = 1;
}
}
for (int i = 0; i < numPartAndBuck; i++) {
newSortNullOrder.add(nullOrder);
}
newSortNullOrder.addAll(sortNullOrder);
String nullOrderStr = "";
for (Integer i : newSortNullOrder) {
if (i.intValue() == 0) {
nullOrderStr += "a";
} else {
nullOrderStr += "z";
}
}
Map<String, ExprNodeDesc> colExprMap = Maps.newHashMap();
ArrayList<ExprNodeDesc> partCols = Lists.newArrayList();
// we will clone here as RS will update bucket column key with its
// corresponding with bucket number and hence their OIs
for (Integer idx : keyColsPosInVal) {
if (idx < 0) {
ExprNodeConstantDesc bucketNumCol =
new ExprNodeConstantDesc(TypeInfoFactory.stringTypeInfo, BUCKET_NUMBER_COL_NAME);
keyCols.add(bucketNumCol);
colExprMap.put(
Utilities.ReduceField.KEY + ".'" + BUCKET_NUMBER_COL_NAME + "'", bucketNumCol);
} else {
keyCols.add(allCols.get(idx).clone());
}
}
ArrayList<ExprNodeDesc> valCols = Lists.newArrayList();
for (int i = 0; i < allCols.size(); i++) {
if (!keyColsPosInVal.contains(i)) {
valCols.add(allCols.get(i).clone());
}
}
for (Integer idx : partitionPositions) {
partCols.add(allCols.get(idx).clone());
}
// in the absence of SORTED BY clause, the sorted dynamic partition insert
// should honor the ordering of records provided by ORDER BY in SELECT statement
ReduceSinkOperator parentRSOp =
OperatorUtils.findSingleOperatorUpstream(parent, ReduceSinkOperator.class);
if (parentRSOp != null && parseCtx.getQueryProperties().hasOuterOrderBy()) {
String parentRSOpOrder = parentRSOp.getConf().getOrder();
String parentRSOpNullOrder = parentRSOp.getConf().getNullOrder();
if (parentRSOpOrder != null && !parentRSOpOrder.isEmpty() && sortPositions.isEmpty()) {
keyCols.addAll(parentRSOp.getConf().getKeyCols());
orderStr += parentRSOpOrder;
nullOrderStr += parentRSOpNullOrder;
}
}
// map _col0 to KEY._col0, etc
Map<String, String> nameMapping = new HashMap<>();
ArrayList<String> keyColNames = Lists.newArrayList();
for (ExprNodeDesc keyCol : keyCols) {
String keyColName = keyCol.getExprString();
keyColNames.add(keyColName);
colExprMap.put(Utilities.ReduceField.KEY + "." + keyColName, keyCol);
nameMapping.put(keyColName, Utilities.ReduceField.KEY + "." + keyColName);
}
ArrayList<String> valColNames = Lists.newArrayList();
for (ExprNodeDesc valCol : valCols) {
String colName = valCol.getExprString();
valColNames.add(colName);
colExprMap.put(Utilities.ReduceField.VALUE + "." + colName, valCol);
nameMapping.put(colName, Utilities.ReduceField.VALUE + "." + colName);
}
// Create Key/Value TableDesc. When the operator plan is split into MR tasks,
// the reduce operator will initialize Extract operator with information
// from Key and Value TableDesc
List<FieldSchema> fields =
PlanUtils.getFieldSchemasFromColumnList(keyCols, keyColNames, 0, "");
TableDesc keyTable = PlanUtils.getReduceKeyTableDesc(fields, orderStr, nullOrderStr);
List<FieldSchema> valFields =
PlanUtils.getFieldSchemasFromColumnList(valCols, valColNames, 0, "");
TableDesc valueTable = PlanUtils.getReduceValueTableDesc(valFields);
List<List<Integer>> distinctColumnIndices = Lists.newArrayList();
// Number of reducers is set to default (-1)
ReduceSinkDesc rsConf =
new ReduceSinkDesc(
keyCols,
keyCols.size(),
valCols,
keyColNames,
distinctColumnIndices,
valColNames,
-1,
partCols,
-1,
keyTable,
valueTable,
writeType);
rsConf.setBucketCols(bucketColumns);
rsConf.setNumBuckets(numBuckets);
ArrayList<ColumnInfo> signature = new ArrayList<>();
for (int index = 0; index < parent.getSchema().getSignature().size(); index++) {
ColumnInfo colInfo = new ColumnInfo(parent.getSchema().getSignature().get(index));
colInfo.setInternalName(nameMapping.get(colInfo.getInternalName()));
signature.add(colInfo);
}
ReduceSinkOperator op =
(ReduceSinkOperator)
OperatorFactory.getAndMakeChild(rsConf, new RowSchema(signature), parent);
op.setColumnExprMap(colExprMap);
return op;
}