tajo-core/src/main/java/org/apache/tajo/engine/planner/physical/DistinctGroupbyThirdAggregationExec.java - tajo - Git at Google

 /**
  * 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.tajo.engine.planner.physical;

 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.apache.tajo.catalog.Column;
 import org.apache.tajo.plan.Target;
 import org.apache.tajo.plan.expr.AggregationFunctionCallEval;
 import org.apache.tajo.plan.function.FunctionContext;
 import org.apache.tajo.plan.logical.DistinctGroupbyNode;
 import org.apache.tajo.plan.logical.GroupbyNode;
 import org.apache.tajo.storage.Tuple;
 import org.apache.tajo.storage.VTuple;
 import org.apache.tajo.worker.TaskAttemptContext;

 import java.io.IOException;
 import java.util.*;

 /**
  *  This class aggregates the output of DistinctGroupbySecondAggregationExec.
  *
  */
 public class DistinctGroupbyThirdAggregationExec extends UnaryPhysicalExec {
   private static Log LOG = LogFactory.getLog(DistinctGroupbyThirdAggregationExec.class);
   private DistinctGroupbyNode plan;

   private boolean finished = false;

   private DistinctFinalAggregator[] aggregators;
   private DistinctFinalAggregator nonDistinctAggr;

   private int resultTupleLength;
   private int numGroupingColumns;

   private int[] resultTupleIndexes;

   private Tuple outTuple;
   private Tuple keyTuple;
   private Tuple prevKeyTuple = null;
   private Tuple prevTuple = null;

   public DistinctGroupbyThirdAggregationExec(TaskAttemptContext context, DistinctGroupbyNode plan, SortExec sortExec)
       throws IOException {
     super(context, plan.getInSchema(), plan.getOutSchema(), sortExec);
     this.plan = plan;
   }

   @Override
   public void init() throws IOException {
     super.init();

     numGroupingColumns = plan.getGroupingColumns().length;
     resultTupleLength = numGroupingColumns;
     keyTuple = new VTuple(numGroupingColumns);

     List<GroupbyNode> groupbyNodes = plan.getSubPlans();

     List<DistinctFinalAggregator> aggregatorList = new ArrayList<>();
     int inTupleIndex = 1 + numGroupingColumns;
     int outTupleIndex = numGroupingColumns;
     int distinctSeq = 0;

     for (GroupbyNode eachGroupby : groupbyNodes) {
       if (eachGroupby.isDistinct()) {
         aggregatorList.add(new DistinctFinalAggregator(distinctSeq, inTupleIndex, outTupleIndex, eachGroupby));
         distinctSeq++;

         Column[] distinctGroupingColumns = eachGroupby.getGroupingColumns();
         inTupleIndex += distinctGroupingColumns.length;
         outTupleIndex += eachGroupby.getAggFunctions().size();
       } else {
         nonDistinctAggr = new DistinctFinalAggregator(-1, inTupleIndex, outTupleIndex, eachGroupby);
         outTupleIndex += eachGroupby.getAggFunctions().size();
       }
       resultTupleLength += eachGroupby.getAggFunctions().size();
     }
     aggregators = aggregatorList.toArray(new DistinctFinalAggregator[aggregatorList.size()]);
     outTuple = new VTuple(resultTupleLength);

     // make output schema mapping index
     resultTupleIndexes = new int[outSchema.size()];
     Map<Column, Integer> groupbyResultTupleIndex = new HashMap<>();
     int resultTupleIndex = 0;
     for (Column eachColumn: plan.getGroupingColumns()) {
       groupbyResultTupleIndex.put(eachColumn, resultTupleIndex);
       resultTupleIndex++;
     }
     for (GroupbyNode eachGroupby : groupbyNodes) {
       Set<Column> groupingColumnSet = new HashSet<>();
       Collections.addAll(groupingColumnSet, eachGroupby.getGroupingColumns());
       for (Target eachTarget: eachGroupby.getTargets()) {
         if (!groupingColumnSet.contains(eachTarget.getNamedColumn())) {
           //aggr function
           groupbyResultTupleIndex.put(eachTarget.getNamedColumn(), resultTupleIndex);
           resultTupleIndex++;
         }
       }
     }

     int index = 0;
     for (Column eachOutputColumn: outSchema.getRootColumns()) {
       // If column is avg aggregation function, outschema's column type is float
       // but groupbyResultTupleIndex's column type is protobuf

       int matchedIndex = -1;
       for (Map.Entry<Column, Integer> entry: groupbyResultTupleIndex.entrySet()) {
         if (entry.getKey().getQualifiedName().equals(eachOutputColumn.getQualifiedName())) {
           matchedIndex = entry.getValue();
           break;
         }
       }
       if (matchedIndex < 0) {
         throw new IOException("Can't find proper output column mapping: " + eachOutputColumn);
       }
       resultTupleIndexes[matchedIndex] = index++;
     }
   }

   @Override
   public Tuple next() throws IOException {
     if (finished) {
       return null;
     }

     while (!context.isStopped()) {
       Tuple tuple = child.next();
       // Last tuple
       if (tuple == null) {
         finished = true;

         if (prevTuple == null) {
           // Empty case
           if (numGroupingColumns == 0) {
             // No grouping column, return null tuple
             return makeEmptyTuple();
           } else {
             return null;
           }
         }

         for (int i = 0; i < numGroupingColumns; i++) {
           outTuple.put(resultTupleIndexes[i], prevTuple.asDatum(i + 1));
         }
         for (DistinctFinalAggregator eachAggr: aggregators) {
           eachAggr.terminate(outTuple);
         }

         return outTuple;
       }

       int distinctSeq = tuple.getInt2(0);
       Tuple keyTuple = getGroupingKeyTuple(tuple);

       // First tuple
       if (prevKeyTuple == null) {
         prevKeyTuple = new VTuple(keyTuple.getValues());
         prevTuple = new VTuple(tuple.getValues());

         aggregators[distinctSeq].merge(tuple);
         continue;
       }

       if (!prevKeyTuple.equals(keyTuple)) {
         // new grouping key
         for (int i = 0; i < numGroupingColumns; i++) {
           outTuple.put(resultTupleIndexes[i], prevTuple.asDatum(i + 1));
         }
         for (DistinctFinalAggregator eachAggr: aggregators) {
           eachAggr.terminate(outTuple);
         }

         prevKeyTuple.put(keyTuple.getValues());
         prevTuple.put(tuple.getValues());

         aggregators[distinctSeq].merge(tuple);
         return outTuple;
       } else {
         prevKeyTuple.put(keyTuple.getValues());
         prevTuple.put(tuple.getValues());
         aggregators[distinctSeq].merge(tuple);
       }
     }

     return null;
   }

   private Tuple makeEmptyTuple() {
     for (DistinctFinalAggregator eachAggr: aggregators) {
       eachAggr.terminateEmpty(outTuple);
     }

     return outTuple;
   }

   private Tuple getGroupingKeyTuple(Tuple tuple) {
     for (int i = 0; i < numGroupingColumns; i++) {
       keyTuple.put(i, tuple.asDatum(i + 1));
     }

     return keyTuple;
   }

   @Override
   public void rescan() throws IOException {
     super.rescan();
     prevKeyTuple = null;
     prevTuple = null;
     finished = false;
   }

   @Override
   public void close() throws IOException {
     super.close();
   }

   class DistinctFinalAggregator {
     private FunctionContext[] functionContexts;
     private List<AggregationFunctionCallEval> aggrFunctions;
     private int seq;
     private int inTupleIndex;
     private int outTupleIndex;
     public DistinctFinalAggregator(int seq, int inTupleIndex, int outTupleIndex, GroupbyNode groupbyNode) {
       this.seq = seq;
       this.inTupleIndex = inTupleIndex;
       this.outTupleIndex = outTupleIndex;

       aggrFunctions = groupbyNode.getAggFunctions();
       if (aggrFunctions != null) {
         for (AggregationFunctionCallEval eachFunction: aggrFunctions) {
           eachFunction.bind(context.getEvalContext(), inSchema);
           eachFunction.setLastPhase();
         }
       }
       newFunctionContext();
     }

     private void newFunctionContext() {
       functionContexts = new FunctionContext[aggrFunctions.size()];
       for (int i = 0; i < aggrFunctions.size(); i++) {
         functionContexts[i] = aggrFunctions.get(i).newContext();
       }
     }

     public void merge(Tuple tuple) {
       for (int i = 0; i < aggrFunctions.size(); i++) {
         aggrFunctions.get(i).merge(functionContexts[i], tuple);
       }

       if (seq == 0 && nonDistinctAggr != null) {
         nonDistinctAggr.merge(tuple);
       }
     }

     public void terminate(Tuple resultTuple) {
       for (int i = 0; i < aggrFunctions.size(); i++) {
         resultTuple.put(resultTupleIndexes[outTupleIndex + i], aggrFunctions.get(i).terminate(functionContexts[i]));
       }
       newFunctionContext();

       if (seq == 0 && nonDistinctAggr != null) {
         nonDistinctAggr.terminate(resultTuple);
       }
     }

     public void terminateEmpty(Tuple resultTuple) {
       newFunctionContext();
       for (int i = 0; i < aggrFunctions.size(); i++) {
         resultTuple.put(resultTupleIndexes[outTupleIndex + i], aggrFunctions.get(i).terminate(functionContexts[i]));
       }
       if (seq == 0 && nonDistinctAggr != null) {
         nonDistinctAggr.terminateEmpty(resultTuple);
       }
     }
   }
 }
	/**
	* 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.tajo.engine.planner.physical;

	import org.apache.commons.logging.Log;
	import org.apache.commons.logging.LogFactory;
	import org.apache.tajo.catalog.Column;
	import org.apache.tajo.plan.Target;
	import org.apache.tajo.plan.expr.AggregationFunctionCallEval;
	import org.apache.tajo.plan.function.FunctionContext;
	import org.apache.tajo.plan.logical.DistinctGroupbyNode;
	import org.apache.tajo.plan.logical.GroupbyNode;
	import org.apache.tajo.storage.Tuple;
	import org.apache.tajo.storage.VTuple;
	import org.apache.tajo.worker.TaskAttemptContext;

	import java.io.IOException;
	import java.util.*;

	/**
	* This class aggregates the output of DistinctGroupbySecondAggregationExec.
	*
	*/
	public class DistinctGroupbyThirdAggregationExec extends UnaryPhysicalExec {
	private static Log LOG = LogFactory.getLog(DistinctGroupbyThirdAggregationExec.class);
	private DistinctGroupbyNode plan;

	private boolean finished = false;

	private DistinctFinalAggregator[] aggregators;
	private DistinctFinalAggregator nonDistinctAggr;

	private int resultTupleLength;
	private int numGroupingColumns;

	private int[] resultTupleIndexes;

	private Tuple outTuple;
	private Tuple keyTuple;
	private Tuple prevKeyTuple = null;
	private Tuple prevTuple = null;

	public DistinctGroupbyThirdAggregationExec(TaskAttemptContext context, DistinctGroupbyNode plan, SortExec sortExec)
	throws IOException {
	super(context, plan.getInSchema(), plan.getOutSchema(), sortExec);
	this.plan = plan;
	}

	@Override
	public void init() throws IOException {
	super.init();

	numGroupingColumns = plan.getGroupingColumns().length;
	resultTupleLength = numGroupingColumns;
	keyTuple = new VTuple(numGroupingColumns);

	List<GroupbyNode> groupbyNodes = plan.getSubPlans();

	List<DistinctFinalAggregator> aggregatorList = new ArrayList<>();
	int inTupleIndex = 1 + numGroupingColumns;
	int outTupleIndex = numGroupingColumns;
	int distinctSeq = 0;

	for (GroupbyNode eachGroupby : groupbyNodes) {
	if (eachGroupby.isDistinct()) {
	aggregatorList.add(new DistinctFinalAggregator(distinctSeq, inTupleIndex, outTupleIndex, eachGroupby));
	distinctSeq++;

	Column[] distinctGroupingColumns = eachGroupby.getGroupingColumns();
	inTupleIndex += distinctGroupingColumns.length;
	outTupleIndex += eachGroupby.getAggFunctions().size();
	} else {
	nonDistinctAggr = new DistinctFinalAggregator(-1, inTupleIndex, outTupleIndex, eachGroupby);
	outTupleIndex += eachGroupby.getAggFunctions().size();
	}
	resultTupleLength += eachGroupby.getAggFunctions().size();
	}
	aggregators = aggregatorList.toArray(new DistinctFinalAggregator[aggregatorList.size()]);
	outTuple = new VTuple(resultTupleLength);

	// make output schema mapping index
	resultTupleIndexes = new int[outSchema.size()];
	Map<Column, Integer> groupbyResultTupleIndex = new HashMap<>();
	int resultTupleIndex = 0;
	for (Column eachColumn: plan.getGroupingColumns()) {
	groupbyResultTupleIndex.put(eachColumn, resultTupleIndex);
	resultTupleIndex++;
	}
	for (GroupbyNode eachGroupby : groupbyNodes) {
	Set<Column> groupingColumnSet = new HashSet<>();
	Collections.addAll(groupingColumnSet, eachGroupby.getGroupingColumns());
	for (Target eachTarget: eachGroupby.getTargets()) {
	if (!groupingColumnSet.contains(eachTarget.getNamedColumn())) {
	//aggr function
	groupbyResultTupleIndex.put(eachTarget.getNamedColumn(), resultTupleIndex);
	resultTupleIndex++;
	}
	}
	}

	int index = 0;
	for (Column eachOutputColumn: outSchema.getRootColumns()) {
	// If column is avg aggregation function, outschema's column type is float
	// but groupbyResultTupleIndex's column type is protobuf

	int matchedIndex = -1;
	for (Map.Entry<Column, Integer> entry: groupbyResultTupleIndex.entrySet()) {
	if (entry.getKey().getQualifiedName().equals(eachOutputColumn.getQualifiedName())) {
	matchedIndex = entry.getValue();
	break;
	}
	}
	if (matchedIndex < 0) {
	throw new IOException("Can't find proper output column mapping: " + eachOutputColumn);
	}
	resultTupleIndexes[matchedIndex] = index++;
	}
	}

	@Override
	public Tuple next() throws IOException {
	if (finished) {
	return null;
	}

	while (!context.isStopped()) {
	Tuple tuple = child.next();
	// Last tuple
	if (tuple == null) {
	finished = true;

	if (prevTuple == null) {
	// Empty case
	if (numGroupingColumns == 0) {
	// No grouping column, return null tuple
	return makeEmptyTuple();
	} else {
	return null;
	}
	}

	for (int i = 0; i < numGroupingColumns; i++) {
	outTuple.put(resultTupleIndexes[i], prevTuple.asDatum(i + 1));
	}
	for (DistinctFinalAggregator eachAggr: aggregators) {
	eachAggr.terminate(outTuple);
	}

	return outTuple;
	}

	int distinctSeq = tuple.getInt2(0);
	Tuple keyTuple = getGroupingKeyTuple(tuple);

	// First tuple
	if (prevKeyTuple == null) {
	prevKeyTuple = new VTuple(keyTuple.getValues());
	prevTuple = new VTuple(tuple.getValues());

	aggregators[distinctSeq].merge(tuple);
	continue;
	}

	if (!prevKeyTuple.equals(keyTuple)) {
	// new grouping key
	for (int i = 0; i < numGroupingColumns; i++) {
	outTuple.put(resultTupleIndexes[i], prevTuple.asDatum(i + 1));
	}
	for (DistinctFinalAggregator eachAggr: aggregators) {
	eachAggr.terminate(outTuple);
	}

	prevKeyTuple.put(keyTuple.getValues());
	prevTuple.put(tuple.getValues());

	aggregators[distinctSeq].merge(tuple);
	return outTuple;
	} else {
	prevKeyTuple.put(keyTuple.getValues());
	prevTuple.put(tuple.getValues());
	aggregators[distinctSeq].merge(tuple);
	}
	}

	return null;
	}

	private Tuple makeEmptyTuple() {
	for (DistinctFinalAggregator eachAggr: aggregators) {
	eachAggr.terminateEmpty(outTuple);
	}

	return outTuple;
	}

	private Tuple getGroupingKeyTuple(Tuple tuple) {
	for (int i = 0; i < numGroupingColumns; i++) {
	keyTuple.put(i, tuple.asDatum(i + 1));
	}

	return keyTuple;
	}

	@Override
	public void rescan() throws IOException {
	super.rescan();
	prevKeyTuple = null;
	prevTuple = null;
	finished = false;
	}

	@Override
	public void close() throws IOException {
	super.close();
	}

	class DistinctFinalAggregator {
	private FunctionContext[] functionContexts;
	private List<AggregationFunctionCallEval> aggrFunctions;
	private int seq;
	private int inTupleIndex;
	private int outTupleIndex;
	public DistinctFinalAggregator(int seq, int inTupleIndex, int outTupleIndex, GroupbyNode groupbyNode) {
	this.seq = seq;
	this.inTupleIndex = inTupleIndex;
	this.outTupleIndex = outTupleIndex;

	aggrFunctions = groupbyNode.getAggFunctions();
	if (aggrFunctions != null) {
	for (AggregationFunctionCallEval eachFunction: aggrFunctions) {
	eachFunction.bind(context.getEvalContext(), inSchema);
	eachFunction.setLastPhase();
	}
	}
	newFunctionContext();
	}

	private void newFunctionContext() {
	functionContexts = new FunctionContext[aggrFunctions.size()];
	for (int i = 0; i < aggrFunctions.size(); i++) {
	functionContexts[i] = aggrFunctions.get(i).newContext();
	}
	}

	public void merge(Tuple tuple) {
	for (int i = 0; i < aggrFunctions.size(); i++) {
	aggrFunctions.get(i).merge(functionContexts[i], tuple);
	}

	if (seq == 0 && nonDistinctAggr != null) {
	nonDistinctAggr.merge(tuple);
	}
	}

	public void terminate(Tuple resultTuple) {
	for (int i = 0; i < aggrFunctions.size(); i++) {
	resultTuple.put(resultTupleIndexes[outTupleIndex + i], aggrFunctions.get(i).terminate(functionContexts[i]));
	}
	newFunctionContext();

	if (seq == 0 && nonDistinctAggr != null) {
	nonDistinctAggr.terminate(resultTuple);
	}
	}

	public void terminateEmpty(Tuple resultTuple) {
	newFunctionContext();
	for (int i = 0; i < aggrFunctions.size(); i++) {
	resultTuple.put(resultTupleIndexes[outTupleIndex + i], aggrFunctions.get(i).terminate(functionContexts[i]));
	}
	if (seq == 0 && nonDistinctAggr != null) {
	nonDistinctAggr.terminateEmpty(resultTuple);
	}
	}
	}
	}