exec/java-exec/src/test/java/org/apache/drill/test/OperatorTestBuilder.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.test;

 import com.google.common.base.Preconditions;
 import org.apache.commons.lang3.mutable.MutableInt;
 import org.apache.drill.exec.physical.base.AbstractBase;
 import org.apache.drill.exec.physical.base.PhysicalOperator;
 import org.apache.drill.exec.physical.impl.BatchCreator;
 import org.apache.drill.exec.physical.impl.MockRecordBatch;
 import org.apache.drill.exec.physical.impl.svremover.Copier;
 import org.apache.drill.exec.physical.impl.svremover.GenericCopier;
 import org.apache.drill.exec.record.CloseableRecordBatch;
 import org.apache.drill.exec.record.RecordBatch;
 import org.apache.drill.exec.record.RecordBatchSizer;
 import org.apache.drill.exec.record.VectorContainer;
 import org.apache.drill.exec.vector.FixedWidthVector;
 import org.apache.drill.exec.vector.ValueVector;
 import org.apache.drill.exec.vector.VariableWidthVector;
 import org.apache.drill.exec.physical.rowSet.DirectRowSet;
 import org.apache.drill.exec.physical.rowSet.RowSet;
 import org.apache.drill.test.rowSet.RowSetComparison;
 import org.junit.Assert;

 import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
 import java.util.stream.Collectors;

 public class OperatorTestBuilder {

   private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(OperatorTestBuilder.class);

   private final List<RowSet> expectedResults = new ArrayList<>();
   private final List<MockRecordBatch> upstreamBatches = new ArrayList<>();
   private PhysicalOpUnitTestBase physicalOpUnitTestBase;

   private PhysicalOperator physicalOperator;
   private long initReservation = AbstractBase.INIT_ALLOCATION;
   private long maxAllocation = AbstractBase.MAX_ALLOCATION;
   private Optional<Integer> expectedNumBatchesOpt = Optional.empty();
   private Optional<Integer> expectedTotalRowsOpt = Optional.empty();
   private boolean combineOutputBatches;
   private boolean unordered;

   public OperatorTestBuilder(PhysicalOpUnitTestBase physicalOpUnitTestBase) {
     this.physicalOpUnitTestBase = physicalOpUnitTestBase;
   }

   @SuppressWarnings("unchecked")
   public void go() throws Exception {
     final List<RowSet> actualResults = new ArrayList<>();
     CloseableRecordBatch testOperator = null;

     try {
       validate();
       int expectedNumBatches = expectedNumBatchesOpt.orElse(expectedResults.size());
       physicalOpUnitTestBase.mockOpContext(physicalOperator, initReservation, maxAllocation);

       final BatchCreator<PhysicalOperator> opCreator = (BatchCreator<PhysicalOperator>) physicalOpUnitTestBase.opCreatorReg.getOperatorCreator(physicalOperator.getClass());
       testOperator = opCreator.getBatch(physicalOpUnitTestBase.fragContext, physicalOperator, (List)upstreamBatches);

       batchIterator: for (int batchIndex = 0;; batchIndex++) {
         final RecordBatch.IterOutcome outcome = testOperator.next();

         switch (outcome) {
           case NONE:
             if (!combineOutputBatches) {
               Assert.assertEquals(expectedNumBatches, batchIndex);
             }
             // We are done iterating over batches. Now we need to compare them.
             break batchIterator;
           case OK_NEW_SCHEMA:
             boolean skip = true;

             try {
               skip = testOperator.getContainer().getRecordCount() == 0;
             } catch (IllegalStateException e) {
               // We should skip this batch in this case. It means no data was included with the okay schema
             } finally {
               if (skip) {
                 batchIndex--;
                 break;
               }
             }
           case OK:
             if (!combineOutputBatches && batchIndex >= expectedNumBatches) {
               testOperator.getContainer().clear();
               Assert.fail("More batches received than expected.");
             } else {
               final boolean hasSelectionVector = testOperator.getSchema().getSelectionVectorMode().hasSelectionVector;
               final VectorContainer container = testOperator.getContainer();

               if (hasSelectionVector) {
                 throw new UnsupportedOperationException("Implement DRILL-6698");
               } else {
                 actualResults.add(DirectRowSet.fromContainer(container));
               }
               break;
             }
           default:
             throw new UnsupportedOperationException("Can't handle this yet");
         }
       }

       int actualTotalRows = actualResults.stream()
         .mapToInt(RowSet::rowCount)
         .reduce(Integer::sum)
         .orElse(0);

       if (expectedResults.isEmpty()) {
         Assert.assertEquals((int) expectedTotalRowsOpt.orElse(0), actualTotalRows);
         // We are done, we don't have any expected result to compare
         return;
       }

       if (combineOutputBatches) {
         final RowSet expectedBatch = expectedResults.get(0);
         final RowSet actualBatch = DirectRowSet.fromSchema(
           physicalOpUnitTestBase.operatorFixture.allocator, actualResults.get(0).container().getSchema());
         final VectorContainer actualBatchContainer = actualBatch.container();
         actualBatchContainer.setRecordCount(0);

         final int numColumns = expectedBatch.schema().size();
         List<MutableInt> totalBytesPerColumn = new ArrayList<>();

         for (int columnIndex = 0; columnIndex < numColumns; columnIndex++) {
           totalBytesPerColumn.add(new MutableInt());
         }

         // Get column sizes for each result batch

         final List<List<RecordBatchSizer.ColumnSize>> columnSizesPerBatch = actualResults.stream().map(rowSet -> {
           switch (rowSet.indirectionType()) {
             case NONE:
               return new RecordBatchSizer(rowSet.container()).columnsList();
             default:
               throw new UnsupportedOperationException("Implement DRILL-6698");
           }
         }).collect(Collectors.toList());

         // Get total bytes per column

         for (List<RecordBatchSizer.ColumnSize> columnSizes: columnSizesPerBatch) {
           for (int columnIndex = 0; columnIndex < numColumns; columnIndex++) {
             final MutableInt totalBytes = totalBytesPerColumn.get(columnIndex);
             final RecordBatchSizer.ColumnSize columnSize = columnSizes.get(columnIndex);
             totalBytes.add(columnSize.getTotalDataSize());
           }
         }

         for (int columnIndex = 0; columnIndex < numColumns; columnIndex++) {
           final ValueVector valueVector = actualBatchContainer
             .getValueVector(columnIndex)
             .getValueVector();

           if (valueVector instanceof FixedWidthVector) {
             ((FixedWidthVector) valueVector).allocateNew(actualTotalRows);
           } else if (valueVector instanceof VariableWidthVector) {
             final MutableInt totalBytes = totalBytesPerColumn.get(columnIndex);
             ((VariableWidthVector) valueVector).allocateNew(totalBytes.getValue(), actualTotalRows);
           } else {
             throw new UnsupportedOperationException();
           }
         }

         try {
           int currentIndex = 0;

           for (RowSet actualRowSet: actualResults) {
             final Copier copier;
             final VectorContainer rowSetContainer = actualRowSet.container();
             rowSetContainer.setRecordCount(actualRowSet.rowCount());

             switch (actualRowSet.indirectionType()) {
               case NONE:
                 copier = new GenericCopier();
                 break;
               default:
                 throw new UnsupportedOperationException("Implement DRILL-6698");
             }

             copier.setup(rowSetContainer, actualBatchContainer);
             copier.appendRecords(currentIndex, actualRowSet.rowCount());
             currentIndex += actualRowSet.rowCount();

             verify(expectedBatch, actualBatch);
           }
         } finally {
           actualBatch.clear();
         }
       } else {
         // Compare expected and actual results
         for (int batchIndex = 0; batchIndex < expectedNumBatches; batchIndex++) {
           final RowSet expectedBatch = expectedResults.get(batchIndex);
           final RowSet actualBatch = actualResults.get(batchIndex);

           verify(expectedBatch, actualBatch);
         }
       }
     } finally {
       // free resources

       if (testOperator != null) {
         testOperator.close();
       }

       actualResults.forEach(rowSet -> rowSet.clear());

       if (expectedResults != null) {
         expectedResults.forEach(rowSet -> rowSet.clear());
       }

       upstreamBatches.forEach(rowSetBatch -> {
         try {
           rowSetBatch.close();
         } catch (Exception e) {
           logger.error("Error while closing RowSetBatch", e);
         }
       });
     }
   }

   private void verify(final RowSet expectedBatch, final RowSet actualBatch) {
     if (unordered) {
       new RowSetComparison(expectedBatch).unorderedVerify(actualBatch);
     } else {
       new RowSetComparison(expectedBatch).verify(actualBatch);
     }
   }

   /**
    * Make sure the inputs are valid.
    */
   private void validate() {
     if (combineOutputBatches) {
       Preconditions.checkArgument(expectedResults.isEmpty() || expectedResults.size() == 1,
         "The number of expected result batches needs to be zero or one when combining output batches");
       Preconditions.checkArgument((expectedResults.isEmpty() && (!expectedNumBatchesOpt.isPresent() && expectedTotalRowsOpt.isPresent())) ||
           (!expectedResults.isEmpty() && (!expectedNumBatchesOpt.isPresent() && !expectedTotalRowsOpt.isPresent())),
         "When definig expectedResults, you cannot define expectedNumBatch or expectedTotalRows and vice versa");
     } else {
       Preconditions.checkArgument((expectedResults.isEmpty() && (expectedNumBatchesOpt.isPresent() || expectedTotalRowsOpt.isPresent())) ||
           (!expectedResults.isEmpty() && (!expectedNumBatchesOpt.isPresent() && !expectedTotalRowsOpt.isPresent())),
         "When definig expectedResults, you cannot define expectedNumBatch or expectedTotalRows and vice versa");
     }
   }

   public OperatorTestBuilder physicalOperator(PhysicalOperator batch) {
     this.physicalOperator = batch;
     return this;
   }

   public OperatorTestBuilder initReservation(long initReservation) {
     this.initReservation = initReservation;
     return this;
   }

   public OperatorTestBuilder maxAllocation(long maxAllocation) {
     this.maxAllocation = maxAllocation;
     return this;
   }

   public OperatorTestBuilder expectedNumBatches(int expectedNumBatches) {
     this.expectedNumBatchesOpt = Optional.of(expectedNumBatches);
     return this;
   }

   public OperatorTestBuilder expectedTotalRows(int expectedTotalRows) {
     this.expectedTotalRowsOpt = Optional.of(expectedTotalRows);
     return this;
   }

   /**
    * Combines all the batches output by the operator into a single batch for comparison.
    * @return This {@link OperatorTestBuilder}.
    */
   public OperatorTestBuilder combineOutputBatches() {
     combineOutputBatches = true;
     return this;
   }

   public OperatorTestBuilder unordered() {
     unordered = true;
     return this;
   }

   public OperatorTestBuilder addUpstreamBatch(final MockRecordBatch mockRecordBatch) {
     Preconditions.checkNotNull(mockRecordBatch);
     upstreamBatches.add(mockRecordBatch);
     return this;
   }

   public OperatorTestBuilder addExpectedResult(final RowSet rowSet) {
     Preconditions.checkNotNull(rowSet);
     expectedResults.add(rowSet);
     return this;
   }
 }
	/*
	* 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.test;

	import com.google.common.base.Preconditions;
	import org.apache.commons.lang3.mutable.MutableInt;
	import org.apache.drill.exec.physical.base.AbstractBase;
	import org.apache.drill.exec.physical.base.PhysicalOperator;
	import org.apache.drill.exec.physical.impl.BatchCreator;
	import org.apache.drill.exec.physical.impl.MockRecordBatch;
	import org.apache.drill.exec.physical.impl.svremover.Copier;
	import org.apache.drill.exec.physical.impl.svremover.GenericCopier;
	import org.apache.drill.exec.record.CloseableRecordBatch;
	import org.apache.drill.exec.record.RecordBatch;
	import org.apache.drill.exec.record.RecordBatchSizer;
	import org.apache.drill.exec.record.VectorContainer;
	import org.apache.drill.exec.vector.FixedWidthVector;
	import org.apache.drill.exec.vector.ValueVector;
	import org.apache.drill.exec.vector.VariableWidthVector;
	import org.apache.drill.exec.physical.rowSet.DirectRowSet;
	import org.apache.drill.exec.physical.rowSet.RowSet;
	import org.apache.drill.test.rowSet.RowSetComparison;
	import org.junit.Assert;

	import java.util.ArrayList;
	import java.util.List;
	import java.util.Optional;
	import java.util.stream.Collectors;

	public class OperatorTestBuilder {

	private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(OperatorTestBuilder.class);

	private final List<RowSet> expectedResults = new ArrayList<>();
	private final List<MockRecordBatch> upstreamBatches = new ArrayList<>();
	private PhysicalOpUnitTestBase physicalOpUnitTestBase;

	private PhysicalOperator physicalOperator;
	private long initReservation = AbstractBase.INIT_ALLOCATION;
	private long maxAllocation = AbstractBase.MAX_ALLOCATION;
	private Optional<Integer> expectedNumBatchesOpt = Optional.empty();
	private Optional<Integer> expectedTotalRowsOpt = Optional.empty();
	private boolean combineOutputBatches;
	private boolean unordered;

	public OperatorTestBuilder(PhysicalOpUnitTestBase physicalOpUnitTestBase) {
	this.physicalOpUnitTestBase = physicalOpUnitTestBase;
	}

	@SuppressWarnings("unchecked")
	public void go() throws Exception {
	final List<RowSet> actualResults = new ArrayList<>();
	CloseableRecordBatch testOperator = null;

	try {
	validate();
	int expectedNumBatches = expectedNumBatchesOpt.orElse(expectedResults.size());
	physicalOpUnitTestBase.mockOpContext(physicalOperator, initReservation, maxAllocation);

	final BatchCreator<PhysicalOperator> opCreator = (BatchCreator<PhysicalOperator>) physicalOpUnitTestBase.opCreatorReg.getOperatorCreator(physicalOperator.getClass());
	testOperator = opCreator.getBatch(physicalOpUnitTestBase.fragContext, physicalOperator, (List)upstreamBatches);

	batchIterator: for (int batchIndex = 0;; batchIndex++) {
	final RecordBatch.IterOutcome outcome = testOperator.next();

	switch (outcome) {
	case NONE:
	if (!combineOutputBatches) {
	Assert.assertEquals(expectedNumBatches, batchIndex);
	}
	// We are done iterating over batches. Now we need to compare them.
	break batchIterator;
	case OK_NEW_SCHEMA:
	boolean skip = true;

	try {
	skip = testOperator.getContainer().getRecordCount() == 0;
	} catch (IllegalStateException e) {
	// We should skip this batch in this case. It means no data was included with the okay schema
	} finally {
	if (skip) {
	batchIndex--;
	break;
	}
	}
	case OK:
	if (!combineOutputBatches && batchIndex >= expectedNumBatches) {
	testOperator.getContainer().clear();
	Assert.fail("More batches received than expected.");
	} else {
	final boolean hasSelectionVector = testOperator.getSchema().getSelectionVectorMode().hasSelectionVector;
	final VectorContainer container = testOperator.getContainer();

	if (hasSelectionVector) {
	throw new UnsupportedOperationException("Implement DRILL-6698");
	} else {
	actualResults.add(DirectRowSet.fromContainer(container));
	}
	break;
	}
	default:
	throw new UnsupportedOperationException("Can't handle this yet");
	}
	}

	int actualTotalRows = actualResults.stream()
	.mapToInt(RowSet::rowCount)
	.reduce(Integer::sum)
	.orElse(0);

	if (expectedResults.isEmpty()) {
	Assert.assertEquals((int) expectedTotalRowsOpt.orElse(0), actualTotalRows);
	// We are done, we don't have any expected result to compare
	return;
	}

	if (combineOutputBatches) {
	final RowSet expectedBatch = expectedResults.get(0);
	final RowSet actualBatch = DirectRowSet.fromSchema(
	physicalOpUnitTestBase.operatorFixture.allocator, actualResults.get(0).container().getSchema());
	final VectorContainer actualBatchContainer = actualBatch.container();
	actualBatchContainer.setRecordCount(0);

	final int numColumns = expectedBatch.schema().size();
	List<MutableInt> totalBytesPerColumn = new ArrayList<>();

	for (int columnIndex = 0; columnIndex < numColumns; columnIndex++) {
	totalBytesPerColumn.add(new MutableInt());
	}

	// Get column sizes for each result batch

	final List<List<RecordBatchSizer.ColumnSize>> columnSizesPerBatch = actualResults.stream().map(rowSet -> {
	switch (rowSet.indirectionType()) {
	case NONE:
	return new RecordBatchSizer(rowSet.container()).columnsList();
	default:
	throw new UnsupportedOperationException("Implement DRILL-6698");
	}
	}).collect(Collectors.toList());

	// Get total bytes per column

	for (List<RecordBatchSizer.ColumnSize> columnSizes: columnSizesPerBatch) {
	for (int columnIndex = 0; columnIndex < numColumns; columnIndex++) {
	final MutableInt totalBytes = totalBytesPerColumn.get(columnIndex);
	final RecordBatchSizer.ColumnSize columnSize = columnSizes.get(columnIndex);
	totalBytes.add(columnSize.getTotalDataSize());
	}
	}

	for (int columnIndex = 0; columnIndex < numColumns; columnIndex++) {
	final ValueVector valueVector = actualBatchContainer
	.getValueVector(columnIndex)
	.getValueVector();

	if (valueVector instanceof FixedWidthVector) {
	((FixedWidthVector) valueVector).allocateNew(actualTotalRows);
	} else if (valueVector instanceof VariableWidthVector) {
	final MutableInt totalBytes = totalBytesPerColumn.get(columnIndex);
	((VariableWidthVector) valueVector).allocateNew(totalBytes.getValue(), actualTotalRows);
	} else {
	throw new UnsupportedOperationException();
	}
	}

	try {
	int currentIndex = 0;

	for (RowSet actualRowSet: actualResults) {
	final Copier copier;
	final VectorContainer rowSetContainer = actualRowSet.container();
	rowSetContainer.setRecordCount(actualRowSet.rowCount());

	switch (actualRowSet.indirectionType()) {
	case NONE:
	copier = new GenericCopier();
	break;
	default:
	throw new UnsupportedOperationException("Implement DRILL-6698");
	}

	copier.setup(rowSetContainer, actualBatchContainer);
	copier.appendRecords(currentIndex, actualRowSet.rowCount());
	currentIndex += actualRowSet.rowCount();

	verify(expectedBatch, actualBatch);
	}
	} finally {
	actualBatch.clear();
	}
	} else {
	// Compare expected and actual results
	for (int batchIndex = 0; batchIndex < expectedNumBatches; batchIndex++) {
	final RowSet expectedBatch = expectedResults.get(batchIndex);
	final RowSet actualBatch = actualResults.get(batchIndex);

	verify(expectedBatch, actualBatch);
	}
	}
	} finally {
	// free resources

	if (testOperator != null) {
	testOperator.close();
	}

	actualResults.forEach(rowSet -> rowSet.clear());

	if (expectedResults != null) {
	expectedResults.forEach(rowSet -> rowSet.clear());
	}

	upstreamBatches.forEach(rowSetBatch -> {
	try {
	rowSetBatch.close();
	} catch (Exception e) {
	logger.error("Error while closing RowSetBatch", e);
	}
	});
	}
	}

	private void verify(final RowSet expectedBatch, final RowSet actualBatch) {
	if (unordered) {
	new RowSetComparison(expectedBatch).unorderedVerify(actualBatch);
	} else {
	new RowSetComparison(expectedBatch).verify(actualBatch);
	}
	}

	/**
	* Make sure the inputs are valid.
	*/
	private void validate() {
	if (combineOutputBatches) {
	Preconditions.checkArgument(expectedResults.isEmpty() \|\| expectedResults.size() == 1,
	"The number of expected result batches needs to be zero or one when combining output batches");
	Preconditions.checkArgument((expectedResults.isEmpty() && (!expectedNumBatchesOpt.isPresent() && expectedTotalRowsOpt.isPresent())) \|\|
	(!expectedResults.isEmpty() && (!expectedNumBatchesOpt.isPresent() && !expectedTotalRowsOpt.isPresent())),
	"When definig expectedResults, you cannot define expectedNumBatch or expectedTotalRows and vice versa");
	} else {
	Preconditions.checkArgument((expectedResults.isEmpty() && (expectedNumBatchesOpt.isPresent() \|\| expectedTotalRowsOpt.isPresent())) \|\|
	(!expectedResults.isEmpty() && (!expectedNumBatchesOpt.isPresent() && !expectedTotalRowsOpt.isPresent())),
	"When definig expectedResults, you cannot define expectedNumBatch or expectedTotalRows and vice versa");
	}
	}

	public OperatorTestBuilder physicalOperator(PhysicalOperator batch) {
	this.physicalOperator = batch;
	return this;
	}

	public OperatorTestBuilder initReservation(long initReservation) {
	this.initReservation = initReservation;
	return this;
	}

	public OperatorTestBuilder maxAllocation(long maxAllocation) {
	this.maxAllocation = maxAllocation;
	return this;
	}

	public OperatorTestBuilder expectedNumBatches(int expectedNumBatches) {
	this.expectedNumBatchesOpt = Optional.of(expectedNumBatches);
	return this;
	}

	public OperatorTestBuilder expectedTotalRows(int expectedTotalRows) {
	this.expectedTotalRowsOpt = Optional.of(expectedTotalRows);
	return this;
	}

	/**
	* Combines all the batches output by the operator into a single batch for comparison.
	* @return This {@link OperatorTestBuilder}.
	*/
	public OperatorTestBuilder combineOutputBatches() {
	combineOutputBatches = true;
	return this;
	}

	public OperatorTestBuilder unordered() {
	unordered = true;
	return this;
	}

	public OperatorTestBuilder addUpstreamBatch(final MockRecordBatch mockRecordBatch) {
	Preconditions.checkNotNull(mockRecordBatch);
	upstreamBatches.add(mockRecordBatch);
	return this;
	}

	public OperatorTestBuilder addExpectedResult(final RowSet rowSet) {
	Preconditions.checkNotNull(rowSet);
	expectedResults.add(rowSet);
	return this;
	}
	}