exec/java-exec/src/main/java/org/apache/drill/exec/physical/impl/unnest/UnnestImpl.java - drill - 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.drill.exec.physical.impl.unnest;

 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import org.apache.drill.exec.exception.SchemaChangeException;
 import org.apache.drill.exec.ops.FragmentContext;
 import org.apache.drill.exec.record.BatchSchema.SelectionVectorMode;
 import org.apache.drill.exec.record.RecordBatch;
 import org.apache.drill.exec.record.TransferPair;
 import org.apache.drill.exec.vector.IntVector;
 import org.apache.drill.exec.vector.SchemaChangeCallBack;
 import org.apache.drill.exec.vector.ValueVector;
 import org.apache.drill.exec.vector.complex.RepeatedValueVector;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

 import java.util.List;

 import static org.apache.drill.exec.record.BatchSchema.SelectionVectorMode.NONE;

 /**
  * Contains the actual unnest operation. Unnest is a simple transfer operation in this implementation.
  * Additionally, unnest produces an implicit rowId column that allows unnest to output batches with many
  * rows of incoming data being unnested in a single call to innerNext(). Downstream blocking operators need
  * to be aware of this rowId column and include the rowId as the sort or group by key.
  * This class follows the pattern of other operators that generate code at runtime. Normally this class
  * would be abstract and have placeholders for doSetup and doEval. Unnest however, doesn't require code
  * generation so we can simply implement the code in a simple class that looks similar to the code gen
  * templates used by other operators but does not implement the doSetup and doEval methods.
  */
 public class UnnestImpl implements Unnest {
   private static final Logger logger = LoggerFactory.getLogger(UnnestImpl.class);

   private ImmutableList<TransferPair> transfers;
   private SelectionVectorMode svMode;
   private RepeatedValueVector fieldToUnnest;
   private RepeatedValueVector.RepeatedAccessor accessor;
   private RecordBatch outgoing;

   private IntVector rowIdVector; // Allocated and owned by the UnnestRecordBatch
   private IntVector.Mutator rowIdVectorMutator;

   /**
    * The output batch limit starts at OUTPUT_ROW_COUNT, but may be decreased
    * if records are found to be large.
    */
   private int outputLimit = ValueVector.MAX_ROW_COUNT;


   /**
    * We maintain three indexes
    *
    *
    *
                 valueIndex  0         1       2   3
                             |- - - - -|- - - -|- -|- - - -|
                             | | | | | | | | | | | | | | | |
                             |- - - - -|- - - -|- -|- - - -|
            innerValueIndex  0 1 2 3 4 0 1 2 3 0 1 0 1 2 3 |
      runningInnerValueIndex 0 1 2 3 4 5 6 7 8 9 ...
    *
    *
    *
    */
   private int valueIndex; // index in the incoming record being processed
   // The index of the array element in the unnest column at row pointed by valueIndex which is currently being
   // processed. It starts at zero and continue until InnerValueCount is reached or the batch limit is reached. It
   // allows for groups to be written across batches if we run out of space. For cases where we have finished
   // a batch on the boundary it will be set to 0
   private int innerValueIndex = 0;
   // The index in the "values" vector of the current value being processed.
   private int runningInnerValueIndex;


   @Override
   public void setUnnestField(RepeatedValueVector unnestField) {
     this.fieldToUnnest = unnestField;
     this.accessor = RepeatedValueVector.RepeatedAccessor.class.cast(unnestField.getAccessor());
   }

   @Override
   public RepeatedValueVector getUnnestField() {
     return fieldToUnnest;
   }

   @Override
   public void setOutputCount(int outputCount) {
     outputLimit = outputCount;
   }

   @Override
   public void setRowIdVector(IntVector v) {
     this.rowIdVector = v;
     this.rowIdVectorMutator = rowIdVector.getMutator();
   }

   @Override
   public final int unnestRecords(final int recordCount) {
     Preconditions.checkArgument(svMode == NONE, "Unnest does not support selection vector inputs.");

     final int initialInnerValueIndex = runningInnerValueIndex;
     int nonEmptyArray = 0;

     outer:
     {
       int outputIndex = 0; // index in the output vector that we are writing to
       final int valueCount = accessor.getValueCount();

       for (; valueIndex < valueCount; valueIndex++) {
         final int innerValueCount = accessor.getInnerValueCountAt(valueIndex);
         logger.trace("Unnest: CurrentRowId: {}, innerValueCount: {}, outputIndex: {},  output limit: {}",
             valueIndex, innerValueCount, outputIndex, outputLimit);

         if (innerValueCount > 0) {
           ++nonEmptyArray;
         }

         for (; innerValueIndex < innerValueCount; innerValueIndex++) {
           // If we've hit the batch size limit, stop and flush what we've got so far.
           if (outputIndex == outputLimit) {
             // Flush this batch.
             break outer;
           }
           try {
             // rowId starts at 1, so the value for rowId is valueIndex+1
             rowIdVectorMutator.setSafe(outputIndex, valueIndex + 1);

           } finally {
             outputIndex++;
             runningInnerValueIndex++;
           }
         }
         innerValueIndex = 0;
       }  // forevery value in the array
     }  // for every incoming record
     final int delta = runningInnerValueIndex - initialInnerValueIndex;
     logger.debug("Unnest: Finished processing current batch. [Details: LastProcessedRowIndex: {}, " +
       "RowsWithNonEmptyArrays: {}, outputIndex: {}, outputLimit: {}, TotalIncomingRecords: {}]",
       valueIndex, nonEmptyArray, delta, outputLimit, accessor.getValueCount());
     final SchemaChangeCallBack callBack = new SchemaChangeCallBack();
     for (TransferPair t : transfers) {
       t.splitAndTransfer(initialInnerValueIndex, delta);

       // Get the corresponding ValueVector in output container and transfer the data
       final ValueVector vectorWithData = t.getTo();
       final ValueVector outputVector = outgoing.getContainer().addOrGet(vectorWithData.getField(), callBack);
       Preconditions.checkState(!callBack.getSchemaChangedAndReset(), "Outgoing container doesn't have "
           + "expected ValueVector of type %s, present in TransferPair of unnest field", vectorWithData.getClass());
       vectorWithData.makeTransferPair(outputVector).transfer();
     }
     return delta;
   }

   @Override
   public final void setup(FragmentContext context, RecordBatch incoming, RecordBatch outgoing,
       List<TransferPair> transfers) throws SchemaChangeException {

     this.svMode = incoming.getSchema().getSelectionVectorMode();
     this.outgoing = outgoing;
     if (svMode == NONE) {
       this.transfers = ImmutableList.copyOf(transfers);
     } else {
       throw new UnsupportedOperationException("Unnest does not support selection vector inputs.");
     }
   }

   @Override
   public void resetGroupIndex() {
     this.valueIndex = 0;
     this.innerValueIndex = 0;
     this.runningInnerValueIndex = 0;
   }

   @Override
   public void close() {
     if (transfers != null) {
       for (TransferPair tp : transfers) {
         tp.getTo().close();
       }
       transfers = null;
     }
   }

   @Override
   public String toString() {
     return "UnnestImpl[svMode=" + svMode
         + ", outputLimit=" + outputLimit
         + ", valueIndex=" + valueIndex
         + ", innerValueIndex=" + innerValueIndex
         + ", runningInnerValueIndex=" + runningInnerValueIndex
         + "]";
   }
 }
	/*
	* 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.drill.exec.physical.impl.unnest;

	import com.google.common.base.Preconditions;
	import com.google.common.collect.ImmutableList;
	import org.apache.drill.exec.exception.SchemaChangeException;
	import org.apache.drill.exec.ops.FragmentContext;
	import org.apache.drill.exec.record.BatchSchema.SelectionVectorMode;
	import org.apache.drill.exec.record.RecordBatch;
	import org.apache.drill.exec.record.TransferPair;
	import org.apache.drill.exec.vector.IntVector;
	import org.apache.drill.exec.vector.SchemaChangeCallBack;
	import org.apache.drill.exec.vector.ValueVector;
	import org.apache.drill.exec.vector.complex.RepeatedValueVector;
	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;

	import java.util.List;

	import static org.apache.drill.exec.record.BatchSchema.SelectionVectorMode.NONE;

	/**
	* Contains the actual unnest operation. Unnest is a simple transfer operation in this implementation.
	* Additionally, unnest produces an implicit rowId column that allows unnest to output batches with many
	* rows of incoming data being unnested in a single call to innerNext(). Downstream blocking operators need
	* to be aware of this rowId column and include the rowId as the sort or group by key.
	* This class follows the pattern of other operators that generate code at runtime. Normally this class
	* would be abstract and have placeholders for doSetup and doEval. Unnest however, doesn't require code
	* generation so we can simply implement the code in a simple class that looks similar to the code gen
	* templates used by other operators but does not implement the doSetup and doEval methods.
	*/
	public class UnnestImpl implements Unnest {
	private static final Logger logger = LoggerFactory.getLogger(UnnestImpl.class);

	private ImmutableList<TransferPair> transfers;
	private SelectionVectorMode svMode;
	private RepeatedValueVector fieldToUnnest;
	private RepeatedValueVector.RepeatedAccessor accessor;
	private RecordBatch outgoing;

	private IntVector rowIdVector; // Allocated and owned by the UnnestRecordBatch
	private IntVector.Mutator rowIdVectorMutator;

	/**
	* The output batch limit starts at OUTPUT_ROW_COUNT, but may be decreased
	* if records are found to be large.
	*/
	private int outputLimit = ValueVector.MAX_ROW_COUNT;



	/**
	* We maintain three indexes
	*
	*
	*
	valueIndex 0 1 2 3
	\|- - - - -\|- - - -\|- -\|- - - -\|
	\| \| \| \| \| \| \| \| \| \| \| \| \| \| \| \|
	\|- - - - -\|- - - -\|- -\|- - - -\|
	innerValueIndex 0 1 2 3 4 0 1 2 3 0 1 0 1 2 3 \|
	runningInnerValueIndex 0 1 2 3 4 5 6 7 8 9 ...
	*
	*
	*
	*/
	private int valueIndex; // index in the incoming record being processed
	// The index of the array element in the unnest column at row pointed by valueIndex which is currently being
	// processed. It starts at zero and continue until InnerValueCount is reached or the batch limit is reached. It
	// allows for groups to be written across batches if we run out of space. For cases where we have finished
	// a batch on the boundary it will be set to 0
	private int innerValueIndex = 0;
	// The index in the "values" vector of the current value being processed.
	private int runningInnerValueIndex;



	@Override
	public void setUnnestField(RepeatedValueVector unnestField) {
	this.fieldToUnnest = unnestField;
	this.accessor = RepeatedValueVector.RepeatedAccessor.class.cast(unnestField.getAccessor());
	}

	@Override
	public RepeatedValueVector getUnnestField() {
	return fieldToUnnest;
	}

	@Override
	public void setOutputCount(int outputCount) {
	outputLimit = outputCount;
	}

	@Override
	public void setRowIdVector(IntVector v) {
	this.rowIdVector = v;
	this.rowIdVectorMutator = rowIdVector.getMutator();
	}

	@Override
	public final int unnestRecords(final int recordCount) {
	Preconditions.checkArgument(svMode == NONE, "Unnest does not support selection vector inputs.");

	final int initialInnerValueIndex = runningInnerValueIndex;
	int nonEmptyArray = 0;

	outer:
	{
	int outputIndex = 0; // index in the output vector that we are writing to
	final int valueCount = accessor.getValueCount();

	for (; valueIndex < valueCount; valueIndex++) {
	final int innerValueCount = accessor.getInnerValueCountAt(valueIndex);
	logger.trace("Unnest: CurrentRowId: {}, innerValueCount: {}, outputIndex: {}, output limit: {}",
	valueIndex, innerValueCount, outputIndex, outputLimit);

	if (innerValueCount > 0) {
	++nonEmptyArray;
	}

	for (; innerValueIndex < innerValueCount; innerValueIndex++) {
	// If we've hit the batch size limit, stop and flush what we've got so far.
	if (outputIndex == outputLimit) {
	// Flush this batch.
	break outer;
	}
	try {
	// rowId starts at 1, so the value for rowId is valueIndex+1
	rowIdVectorMutator.setSafe(outputIndex, valueIndex + 1);

	} finally {
	outputIndex++;
	runningInnerValueIndex++;
	}
	}
	innerValueIndex = 0;
	} // forevery value in the array
	} // for every incoming record
	final int delta = runningInnerValueIndex - initialInnerValueIndex;
	logger.debug("Unnest: Finished processing current batch. [Details: LastProcessedRowIndex: {}, " +
	"RowsWithNonEmptyArrays: {}, outputIndex: {}, outputLimit: {}, TotalIncomingRecords: {}]",
	valueIndex, nonEmptyArray, delta, outputLimit, accessor.getValueCount());
	final SchemaChangeCallBack callBack = new SchemaChangeCallBack();
	for (TransferPair t : transfers) {
	t.splitAndTransfer(initialInnerValueIndex, delta);

	// Get the corresponding ValueVector in output container and transfer the data
	final ValueVector vectorWithData = t.getTo();
	final ValueVector outputVector = outgoing.getContainer().addOrGet(vectorWithData.getField(), callBack);
	Preconditions.checkState(!callBack.getSchemaChangedAndReset(), "Outgoing container doesn't have "
	+ "expected ValueVector of type %s, present in TransferPair of unnest field", vectorWithData.getClass());
	vectorWithData.makeTransferPair(outputVector).transfer();
	}
	return delta;
	}

	@Override
	public final void setup(FragmentContext context, RecordBatch incoming, RecordBatch outgoing,
	List<TransferPair> transfers) throws SchemaChangeException {

	this.svMode = incoming.getSchema().getSelectionVectorMode();
	this.outgoing = outgoing;
	if (svMode == NONE) {
	this.transfers = ImmutableList.copyOf(transfers);
	} else {
	throw new UnsupportedOperationException("Unnest does not support selection vector inputs.");
	}
	}

	@Override
	public void resetGroupIndex() {
	this.valueIndex = 0;
	this.innerValueIndex = 0;
	this.runningInnerValueIndex = 0;
	}

	@Override
	public void close() {
	if (transfers != null) {
	for (TransferPair tp : transfers) {
	tp.getTo().close();
	}
	transfers = null;
	}
	}

	@Override
	public String toString() {
	return "UnnestImpl[svMode=" + svMode
	+ ", outputLimit=" + outputLimit
	+ ", valueIndex=" + valueIndex
	+ ", innerValueIndex=" + innerValueIndex
	+ ", runningInnerValueIndex=" + runningInnerValueIndex
	+ "]";
	}
	}