tajo-core/src/main/java/org/apache/tajo/engine/planner/physical/WindowAggExec.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 com.google.common.collect.Lists;
 import org.apache.tajo.catalog.Column;
 import org.apache.tajo.catalog.Schema;
 import org.apache.tajo.catalog.SortSpec;
 import org.apache.tajo.datum.Datum;
 import org.apache.tajo.plan.expr.EvalNode;
 import org.apache.tajo.plan.expr.WindowFunctionEval;
 import org.apache.tajo.plan.function.FunctionContext;
 import org.apache.tajo.plan.logical.WindowAggNode;
 import org.apache.tajo.plan.logical.WindowSpec;
 import org.apache.tajo.storage.BaseTupleComparator;
 import org.apache.tajo.storage.Tuple;
 import org.apache.tajo.storage.TupleComparator;
 import org.apache.tajo.storage.VTuple;
 import org.apache.tajo.worker.TaskAttemptContext;

 import java.io.IOException;
 import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;

 /**
  * The sort-based window aggregation operator
  */
 public class WindowAggExec extends UnaryPhysicalExec {
   // plan information
   protected final int outputColumnNum;
   protected final int nonFunctionColumnNum;
   protected final int nonFunctionColumns[];

   protected final int functionNum;
   protected final WindowFunctionEval functions[];

   protected Schema schemaForOrderBy;
   protected int sortKeyColumns[];
   protected final boolean hasPartitionKeys;
   protected final int partitionKeyNum;
   protected final int partitionKeyIds[];

   // for evaluation
   protected FunctionContext contexts [];
   protected Tuple lastKey = null;
   protected boolean noMoreTuples = false;
   private boolean [] orderedFuncFlags;
   private boolean [] aggFuncFlags;
   private boolean [] windowFuncFlags;
   private boolean [] endUnboundedFollowingFlags;
   private boolean [] endCurrentRowFlags;
   private Tuple currentKey;
   private Tuple evaluatedTuple;

   // operator state
   enum WindowState {
     NEW_WINDOW,
     ACCUMULATING_WINDOW,
     EVALUATION,
     RETRIEVING_FROM_WINDOW,
     END_OF_TUPLE
   }

   // Transient state
   boolean firstTime = true;
   TupleList evaluatedTuples = null;
   TupleList accumulatedInTuples = null;
   TupleList nextAccumulatedInTuples = null;
   WindowState state = WindowState.NEW_WINDOW;
   Iterator<Tuple> tupleInFrameIterator = null;

   public WindowAggExec(TaskAttemptContext context, WindowAggNode plan, PhysicalExec child) throws IOException {
     super(context, plan.getInSchema(), plan.getOutSchema(), child);

     if (plan.hasPartitionKeys()) {
       final Column[] keyColumns = plan.getPartitionKeys();
       partitionKeyNum = keyColumns.length;
       partitionKeyIds = new int[partitionKeyNum];
       Column col;
       for (int idx = 0; idx < plan.getPartitionKeys().length; idx++) {
         col = keyColumns[idx];
         partitionKeyIds[idx] = inSchema.getColumnId(col.getQualifiedName());
       }
       hasPartitionKeys = true;
     } else {
       partitionKeyNum = 0;
       partitionKeyIds = null;
       hasPartitionKeys = false;
     }

     currentKey = new VTuple(partitionKeyNum);

     if (plan.hasAggFunctions()) {
       functions = plan.getWindowFunctions();
       functionNum = functions.length;

       orderedFuncFlags = new boolean[functions.length];
       windowFuncFlags = new boolean[functions.length];
       aggFuncFlags = new boolean[functions.length];

       endUnboundedFollowingFlags = new boolean[functions.length];
       endCurrentRowFlags = new boolean[functions.length];

       List<Column> additionalSortKeyColumns = Lists.newArrayList();
       Schema rewrittenSchema = new Schema(outSchema);
       for (int i = 0; i < functions.length; i++) {
         WindowSpec.WindowEndBound endBound = functions[i].getWindowFrame().getEndBound();
         switch (endBound.getBoundType()) {
         case CURRENT_ROW:
           endCurrentRowFlags[i] = true; break;
         case UNBOUNDED_FOLLOWING:
           endUnboundedFollowingFlags[i] = true; break;
         default:
         }

         switch (functions[i].getFuncDesc().getFuncType()) {
         case AGGREGATION:
         case DISTINCT_AGGREGATION:
           aggFuncFlags[i] = true; break;
         case WINDOW:
           windowFuncFlags[i] = true; break;
         default:
         }

         if (functions[i].hasSortSpecs()) {
           orderedFuncFlags[i] = true;

           for (SortSpec sortSpec : functions[i].getSortSpecs()) {
             if (!rewrittenSchema.contains(sortSpec.getSortKey())) {
               // check if additionalSortKeyColumns already has that sort key
               boolean newKey = true;
               for (Column c : additionalSortKeyColumns) {
                 if (c.equals(sortSpec.getSortKey())) {
                   newKey = false;
                 }
               }
               if (newKey) {
                 additionalSortKeyColumns.add(sortSpec.getSortKey());
               }
             }
           }
         }
       }

       sortKeyColumns = new int[additionalSortKeyColumns.size()];
       schemaForOrderBy = new Schema(outSchema);
       for (int i = 0; i < additionalSortKeyColumns.size(); i++) {
         sortKeyColumns[i] = i;
         schemaForOrderBy.addColumn(additionalSortKeyColumns.get(i));
       }
     } else {
       functions = new WindowFunctionEval[0];
       functionNum = 0;
       schemaForOrderBy = outSchema;
     }

     evaluatedTuple = new VTuple(schemaForOrderBy.size());
     nonFunctionColumnNum = plan.getTargets().length - functionNum;
     nonFunctionColumns = new int[nonFunctionColumnNum];
     for (int idx = 0; idx < plan.getTargets().length - functionNum; idx++) {
       nonFunctionColumns[idx] = inSchema.getColumnId(plan.getTargets()[idx].getCanonicalName());
     }

     outputColumnNum = nonFunctionColumnNum + functionNum;
   }

   @Override
   public void init() throws IOException {
     super.init();
     for (EvalNode functionEval : functions) {
       functionEval.bind(context.getEvalContext(), inSchema);
     }
   }

   private void transition(WindowState state) {
     this.state = state;
   }

   @Override
   public Tuple next() throws IOException {
     Tuple readTuple = null;

     while(!context.isStopped() && state != WindowState.END_OF_TUPLE) {

       if (state == WindowState.NEW_WINDOW) {
         initWindow();
         transition(WindowState.ACCUMULATING_WINDOW);
       }

       if (state != WindowState.RETRIEVING_FROM_WINDOW) { // read an input tuple and build a partition key
         readTuple = child.next();

         if (readTuple == null) { // the end of tuple
           noMoreTuples = true;
           transition(WindowState.EVALUATION);
         }

         if (readTuple != null && hasPartitionKeys) { // get a key tuple
           for (int i = 0; i < partitionKeyIds.length; i++) {
             currentKey.put(i, readTuple.asDatum(partitionKeyIds[i]));
           }
         }
       }

       if (state == WindowState.ACCUMULATING_WINDOW) {
         accumulatingWindow(currentKey, readTuple);
       }

       if (state == WindowState.EVALUATION) {
         evaluationWindowFrame();

         tupleInFrameIterator = evaluatedTuples.iterator();
         transition(WindowState.RETRIEVING_FROM_WINDOW);
       }

       if (state == WindowState.RETRIEVING_FROM_WINDOW) {
         if (tupleInFrameIterator.hasNext()) {
           return tupleInFrameIterator.next();
         } else {
           finalizeWindow();
         }
       }
     }

     return null;
   }

   private void initWindow() {
     if (firstTime) {
       accumulatedInTuples = new TupleList();

       contexts = new FunctionContext[functionNum];
       for(int evalIdx = 0; evalIdx < functionNum; evalIdx++) {
         contexts[evalIdx] = functions[evalIdx].newContext();
       }
       firstTime = false;
     }
   }

   private void accumulatingWindow(Tuple currentKey, Tuple inTuple) {

     if (lastKey == null || lastKey.equals(currentKey)) { // if the current key is same to the previous key
       accumulatedInTuples.add(inTuple);
       if (lastKey == null) {
         lastKey = new VTuple(currentKey.size());
       }
     } else {
       // if the current key is different from the previous key,
       // the current key belongs to the next window frame. preaccumulatingNextWindow() will
       // aggregate the current key for next window frame.
       preAccumulatingNextWindow(inTuple);
       transition(WindowState.EVALUATION);
     }

     lastKey.put(currentKey.getValues());
   }

   private void preAccumulatingNextWindow(Tuple inTuple) {
     nextAccumulatedInTuples = new TupleList();
     nextAccumulatedInTuples.add(inTuple);
   }

   private void evaluationWindowFrame() {
     TupleComparator comp;

     evaluatedTuples = new TupleList();

     for (int i = 0; i <accumulatedInTuples.size(); i++) {
       Tuple inTuple = accumulatedInTuples.get(i);

       for (int c = 0; c < nonFunctionColumnNum; c++) {
         evaluatedTuple.put(c, inTuple.asDatum(nonFunctionColumns[c]));
       }
       for (int c = 0; c < sortKeyColumns.length; c++) {
         evaluatedTuple.put(outputColumnNum + c, inTuple.asDatum(sortKeyColumns[c]));
       }

       evaluatedTuples.add(evaluatedTuple);
     }

     for (int idx = 0; idx < functions.length; idx++) {
       if (orderedFuncFlags[idx]) {
         comp = new BaseTupleComparator(inSchema, functions[idx].getSortSpecs());
         Collections.sort(accumulatedInTuples, comp);
         comp = new BaseTupleComparator(schemaForOrderBy, functions[idx].getSortSpecs());
         Collections.sort(evaluatedTuples, comp);
       }

       for (int i = 0; i < accumulatedInTuples.size(); i++) {
         Tuple inTuple = accumulatedInTuples.get(i);
         Tuple outTuple = evaluatedTuples.get(i);

         functions[idx].merge(contexts[idx], inTuple);

         if (windowFuncFlags[idx]) {
           Datum result = functions[idx].terminate(contexts[idx]);
           outTuple.put(nonFunctionColumnNum + idx, result);
         }
       }

       if (aggFuncFlags[idx]) {
         for (int i = 0; i < evaluatedTuples.size(); i++) {
           Datum result = functions[idx].terminate(contexts[idx]);
           Tuple outTuple = evaluatedTuples.get(i);
           outTuple.put(nonFunctionColumnNum + idx, result);
         }
       }
     }
   }

   private void finalizeWindow() {
     evaluatedTuples.clear();
     accumulatedInTuples.clear();

     if (noMoreTuples) {
       transition(WindowState.END_OF_TUPLE);
     } else {
       accumulatedInTuples = nextAccumulatedInTuples;

       contexts = new FunctionContext[functionNum];
       for(int evalIdx = 0; evalIdx < functionNum; evalIdx++) {
         contexts[evalIdx] = functions[evalIdx].newContext();
       }
       transition(WindowState.NEW_WINDOW);
     }
   }

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

     lastKey = null;
     noMoreTuples = false;
   }

   @Override
   public void close() throws IOException {
     if (evaluatedTuples != null) {
       evaluatedTuples.clear();
       evaluatedTuples = null;
     }
     if (accumulatedInTuples != null) {
       accumulatedInTuples.clear();
       accumulatedInTuples = null;
     }
     if (nextAccumulatedInTuples != null) {
       nextAccumulatedInTuples.clear();
       nextAccumulatedInTuples = null;
     }

     super.close();
   }
 }
	/**
	* 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 com.google.common.collect.Lists;
	import org.apache.tajo.catalog.Column;
	import org.apache.tajo.catalog.Schema;
	import org.apache.tajo.catalog.SortSpec;
	import org.apache.tajo.datum.Datum;
	import org.apache.tajo.plan.expr.EvalNode;
	import org.apache.tajo.plan.expr.WindowFunctionEval;
	import org.apache.tajo.plan.function.FunctionContext;
	import org.apache.tajo.plan.logical.WindowAggNode;
	import org.apache.tajo.plan.logical.WindowSpec;
	import org.apache.tajo.storage.BaseTupleComparator;
	import org.apache.tajo.storage.Tuple;
	import org.apache.tajo.storage.TupleComparator;
	import org.apache.tajo.storage.VTuple;
	import org.apache.tajo.worker.TaskAttemptContext;

	import java.io.IOException;
	import java.util.Collections;
	import java.util.Iterator;
	import java.util.List;

	/**
	* The sort-based window aggregation operator
	*/
	public class WindowAggExec extends UnaryPhysicalExec {
	// plan information
	protected final int outputColumnNum;
	protected final int nonFunctionColumnNum;
	protected final int nonFunctionColumns[];

	protected final int functionNum;
	protected final WindowFunctionEval functions[];

	protected Schema schemaForOrderBy;
	protected int sortKeyColumns[];
	protected final boolean hasPartitionKeys;
	protected final int partitionKeyNum;
	protected final int partitionKeyIds[];

	// for evaluation
	protected FunctionContext contexts [];
	protected Tuple lastKey = null;
	protected boolean noMoreTuples = false;
	private boolean [] orderedFuncFlags;
	private boolean [] aggFuncFlags;
	private boolean [] windowFuncFlags;
	private boolean [] endUnboundedFollowingFlags;
	private boolean [] endCurrentRowFlags;
	private Tuple currentKey;
	private Tuple evaluatedTuple;

	// operator state
	enum WindowState {
	NEW_WINDOW,
	ACCUMULATING_WINDOW,
	EVALUATION,
	RETRIEVING_FROM_WINDOW,
	END_OF_TUPLE
	}

	// Transient state
	boolean firstTime = true;
	TupleList evaluatedTuples = null;
	TupleList accumulatedInTuples = null;
	TupleList nextAccumulatedInTuples = null;
	WindowState state = WindowState.NEW_WINDOW;
	Iterator<Tuple> tupleInFrameIterator = null;

	public WindowAggExec(TaskAttemptContext context, WindowAggNode plan, PhysicalExec child) throws IOException {
	super(context, plan.getInSchema(), plan.getOutSchema(), child);

	if (plan.hasPartitionKeys()) {
	final Column[] keyColumns = plan.getPartitionKeys();
	partitionKeyNum = keyColumns.length;
	partitionKeyIds = new int[partitionKeyNum];
	Column col;
	for (int idx = 0; idx < plan.getPartitionKeys().length; idx++) {
	col = keyColumns[idx];
	partitionKeyIds[idx] = inSchema.getColumnId(col.getQualifiedName());
	}
	hasPartitionKeys = true;
	} else {
	partitionKeyNum = 0;
	partitionKeyIds = null;
	hasPartitionKeys = false;
	}

	currentKey = new VTuple(partitionKeyNum);

	if (plan.hasAggFunctions()) {
	functions = plan.getWindowFunctions();
	functionNum = functions.length;

	orderedFuncFlags = new boolean[functions.length];
	windowFuncFlags = new boolean[functions.length];
	aggFuncFlags = new boolean[functions.length];

	endUnboundedFollowingFlags = new boolean[functions.length];
	endCurrentRowFlags = new boolean[functions.length];

	List<Column> additionalSortKeyColumns = Lists.newArrayList();
	Schema rewrittenSchema = new Schema(outSchema);
	for (int i = 0; i < functions.length; i++) {
	WindowSpec.WindowEndBound endBound = functions[i].getWindowFrame().getEndBound();
	switch (endBound.getBoundType()) {
	case CURRENT_ROW:
	endCurrentRowFlags[i] = true; break;
	case UNBOUNDED_FOLLOWING:
	endUnboundedFollowingFlags[i] = true; break;
	default:
	}

	switch (functions[i].getFuncDesc().getFuncType()) {
	case AGGREGATION:
	case DISTINCT_AGGREGATION:
	aggFuncFlags[i] = true; break;
	case WINDOW:
	windowFuncFlags[i] = true; break;
	default:
	}

	if (functions[i].hasSortSpecs()) {
	orderedFuncFlags[i] = true;

	for (SortSpec sortSpec : functions[i].getSortSpecs()) {
	if (!rewrittenSchema.contains(sortSpec.getSortKey())) {
	// check if additionalSortKeyColumns already has that sort key
	boolean newKey = true;
	for (Column c : additionalSortKeyColumns) {
	if (c.equals(sortSpec.getSortKey())) {
	newKey = false;
	}
	}
	if (newKey) {
	additionalSortKeyColumns.add(sortSpec.getSortKey());
	}
	}
	}
	}
	}

	sortKeyColumns = new int[additionalSortKeyColumns.size()];
	schemaForOrderBy = new Schema(outSchema);
	for (int i = 0; i < additionalSortKeyColumns.size(); i++) {
	sortKeyColumns[i] = i;
	schemaForOrderBy.addColumn(additionalSortKeyColumns.get(i));
	}
	} else {
	functions = new WindowFunctionEval[0];
	functionNum = 0;
	schemaForOrderBy = outSchema;
	}

	evaluatedTuple = new VTuple(schemaForOrderBy.size());
	nonFunctionColumnNum = plan.getTargets().length - functionNum;
	nonFunctionColumns = new int[nonFunctionColumnNum];
	for (int idx = 0; idx < plan.getTargets().length - functionNum; idx++) {
	nonFunctionColumns[idx] = inSchema.getColumnId(plan.getTargets()[idx].getCanonicalName());
	}

	outputColumnNum = nonFunctionColumnNum + functionNum;
	}

	@Override
	public void init() throws IOException {
	super.init();
	for (EvalNode functionEval : functions) {
	functionEval.bind(context.getEvalContext(), inSchema);
	}
	}

	private void transition(WindowState state) {
	this.state = state;
	}

	@Override
	public Tuple next() throws IOException {
	Tuple readTuple = null;

	while(!context.isStopped() && state != WindowState.END_OF_TUPLE) {

	if (state == WindowState.NEW_WINDOW) {
	initWindow();
	transition(WindowState.ACCUMULATING_WINDOW);
	}

	if (state != WindowState.RETRIEVING_FROM_WINDOW) { // read an input tuple and build a partition key
	readTuple = child.next();

	if (readTuple == null) { // the end of tuple
	noMoreTuples = true;
	transition(WindowState.EVALUATION);
	}

	if (readTuple != null && hasPartitionKeys) { // get a key tuple
	for (int i = 0; i < partitionKeyIds.length; i++) {
	currentKey.put(i, readTuple.asDatum(partitionKeyIds[i]));
	}
	}
	}

	if (state == WindowState.ACCUMULATING_WINDOW) {
	accumulatingWindow(currentKey, readTuple);
	}

	if (state == WindowState.EVALUATION) {
	evaluationWindowFrame();

	tupleInFrameIterator = evaluatedTuples.iterator();
	transition(WindowState.RETRIEVING_FROM_WINDOW);
	}

	if (state == WindowState.RETRIEVING_FROM_WINDOW) {
	if (tupleInFrameIterator.hasNext()) {
	return tupleInFrameIterator.next();
	} else {
	finalizeWindow();
	}
	}
	}

	return null;
	}

	private void initWindow() {
	if (firstTime) {
	accumulatedInTuples = new TupleList();

	contexts = new FunctionContext[functionNum];
	for(int evalIdx = 0; evalIdx < functionNum; evalIdx++) {
	contexts[evalIdx] = functions[evalIdx].newContext();
	}
	firstTime = false;
	}
	}

	private void accumulatingWindow(Tuple currentKey, Tuple inTuple) {

	if (lastKey == null \|\| lastKey.equals(currentKey)) { // if the current key is same to the previous key
	accumulatedInTuples.add(inTuple);
	if (lastKey == null) {
	lastKey = new VTuple(currentKey.size());
	}
	} else {
	// if the current key is different from the previous key,
	// the current key belongs to the next window frame. preaccumulatingNextWindow() will
	// aggregate the current key for next window frame.
	preAccumulatingNextWindow(inTuple);
	transition(WindowState.EVALUATION);
	}

	lastKey.put(currentKey.getValues());
	}

	private void preAccumulatingNextWindow(Tuple inTuple) {
	nextAccumulatedInTuples = new TupleList();
	nextAccumulatedInTuples.add(inTuple);
	}

	private void evaluationWindowFrame() {
	TupleComparator comp;

	evaluatedTuples = new TupleList();

	for (int i = 0; i <accumulatedInTuples.size(); i++) {
	Tuple inTuple = accumulatedInTuples.get(i);

	for (int c = 0; c < nonFunctionColumnNum; c++) {
	evaluatedTuple.put(c, inTuple.asDatum(nonFunctionColumns[c]));
	}
	for (int c = 0; c < sortKeyColumns.length; c++) {
	evaluatedTuple.put(outputColumnNum + c, inTuple.asDatum(sortKeyColumns[c]));
	}

	evaluatedTuples.add(evaluatedTuple);
	}

	for (int idx = 0; idx < functions.length; idx++) {
	if (orderedFuncFlags[idx]) {
	comp = new BaseTupleComparator(inSchema, functions[idx].getSortSpecs());
	Collections.sort(accumulatedInTuples, comp);
	comp = new BaseTupleComparator(schemaForOrderBy, functions[idx].getSortSpecs());
	Collections.sort(evaluatedTuples, comp);
	}

	for (int i = 0; i < accumulatedInTuples.size(); i++) {
	Tuple inTuple = accumulatedInTuples.get(i);
	Tuple outTuple = evaluatedTuples.get(i);

	functions[idx].merge(contexts[idx], inTuple);

	if (windowFuncFlags[idx]) {
	Datum result = functions[idx].terminate(contexts[idx]);
	outTuple.put(nonFunctionColumnNum + idx, result);
	}
	}

	if (aggFuncFlags[idx]) {
	for (int i = 0; i < evaluatedTuples.size(); i++) {
	Datum result = functions[idx].terminate(contexts[idx]);
	Tuple outTuple = evaluatedTuples.get(i);
	outTuple.put(nonFunctionColumnNum + idx, result);
	}
	}
	}
	}

	private void finalizeWindow() {
	evaluatedTuples.clear();
	accumulatedInTuples.clear();

	if (noMoreTuples) {
	transition(WindowState.END_OF_TUPLE);
	} else {
	accumulatedInTuples = nextAccumulatedInTuples;

	contexts = new FunctionContext[functionNum];
	for(int evalIdx = 0; evalIdx < functionNum; evalIdx++) {
	contexts[evalIdx] = functions[evalIdx].newContext();
	}
	transition(WindowState.NEW_WINDOW);
	}
	}

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

	lastKey = null;
	noMoreTuples = false;
	}

	@Override
	public void close() throws IOException {
	if (evaluatedTuples != null) {
	evaluatedTuples.clear();
	evaluatedTuples = null;
	}
	if (accumulatedInTuples != null) {
	accumulatedInTuples.clear();
	accumulatedInTuples = null;
	}
	if (nextAccumulatedInTuples != null) {
	nextAccumulatedInTuples.clear();
	nextAccumulatedInTuples = null;
	}

	super.close();
	}
	}