| /* |
| * 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.collect.Lists; |
| import org.apache.calcite.rel.core.JoinRelType; |
| import org.apache.drill.categories.OperatorTest; |
| import org.apache.drill.common.exceptions.DrillException; |
| import org.apache.drill.common.exceptions.UserException; |
| import org.apache.drill.common.expression.SchemaPath; |
| import org.apache.drill.common.types.TypeProtos; |
| import org.apache.drill.exec.ExecConstants; |
| import org.apache.drill.exec.ops.OperatorContext; |
| import org.apache.drill.exec.physical.base.LateralContract; |
| import org.apache.drill.exec.physical.base.PhysicalOperator; |
| import org.apache.drill.exec.physical.config.LateralJoinPOP; |
| import org.apache.drill.exec.physical.config.Project; |
| import org.apache.drill.exec.physical.config.UnnestPOP; |
| import org.apache.drill.exec.physical.impl.MockRecordBatch; |
| import org.apache.drill.exec.physical.impl.join.LateralJoinBatch; |
| import org.apache.drill.exec.physical.impl.project.ProjectRecordBatch; |
| import org.apache.drill.exec.physical.impl.sort.RecordBatchData; |
| import org.apache.drill.exec.planner.common.DrillLateralJoinRelBase; |
| import org.apache.drill.exec.planner.common.DrillUnnestRelBase; |
| import org.apache.drill.exec.planner.logical.DrillLogicalTestUtils; |
| import org.apache.drill.exec.record.RecordBatch; |
| import org.apache.drill.exec.record.VectorContainer; |
| import org.apache.drill.exec.record.metadata.SchemaBuilder; |
| import org.apache.drill.exec.record.metadata.TupleMetadata; |
| import org.apache.drill.exec.store.mock.MockStorePOP; |
| import org.apache.drill.exec.vector.ValueVector; |
| import org.apache.drill.exec.vector.VarCharVector; |
| import org.apache.drill.exec.vector.complex.MapVector; |
| import org.apache.drill.test.SubOperatorTest; |
| import org.apache.drill.exec.physical.rowSet.RowSet; |
| import org.apache.drill.exec.physical.rowSet.RowSetBuilder; |
| import org.junit.AfterClass; |
| import org.junit.BeforeClass; |
| import org.junit.Test; |
| import org.junit.experimental.categories.Category; |
| |
| import java.util.ArrayList; |
| import java.util.List; |
| |
| import static org.junit.Assert.fail; |
| |
| @Category(OperatorTest.class) |
| public class TestUnnestWithLateralCorrectness extends SubOperatorTest { |
| |
| // Operator Context for mock batch |
| public static OperatorContext operatorContext; |
| |
| public static PhysicalOperator mockPopConfig; |
| public static LateralJoinPOP ljPopConfig; |
| |
| @BeforeClass public static void setUpBeforeClass() throws Exception { |
| mockPopConfig = new MockStorePOP(null); |
| ljPopConfig = new LateralJoinPOP(null, null, JoinRelType.INNER, DrillLateralJoinRelBase.IMPLICIT_COLUMN, Lists.newArrayList()); |
| operatorContext = fixture.newOperatorContext(mockPopConfig); |
| } |
| |
| @AfterClass public static void tearDownAfterClass() throws Exception { |
| operatorContext.close(); |
| } |
| |
| @Test |
| public void testUnnestFixedWidthColumn() { |
| |
| Object[][] data = { |
| { new int[] {1, 2}, |
| new int[] {3, 4, 5}}, |
| { new int[] {6, 7, 8, 9}, |
| new int[] {10, 11, 12, 13, 14}} |
| }; |
| |
| // Create input schema |
| TupleMetadata incomingSchema = |
| new SchemaBuilder() |
| .add("rowNumber", TypeProtos.MinorType.INT) |
| .addArray("unnestColumn", TypeProtos.MinorType.INT) |
| .buildSchema(); |
| TupleMetadata[] incomingSchemas = { incomingSchema, incomingSchema }; |
| |
| Integer[][][] baseline = { |
| { |
| {1, 1, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4, 4}, //rowNum |
| {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14} // unnestColumn_flat |
| } |
| }; |
| |
| RecordBatch.IterOutcome[] iterOutcomes = {RecordBatch.IterOutcome.OK_NEW_SCHEMA, RecordBatch.IterOutcome.OK}; |
| |
| try { |
| testUnnest(incomingSchemas, iterOutcomes, data, baseline, false); |
| } catch (Exception e) { |
| fail("Failed due to exception: " + e.getMessage()); |
| } |
| } |
| |
| @Test |
| public void testUnnestVarWidthColumn() { |
| |
| Object[][] data = { |
| { new String[] {"", "zero"}, |
| new String[] {"one", "two", "three"}}, |
| { new String[] {"four", "five", "six", "seven"}, |
| new String[] {"eight", "nine", "ten", "eleven", "twelve"}} |
| }; |
| |
| // Create input schema |
| TupleMetadata incomingSchema = new SchemaBuilder() |
| .add("someColumn", TypeProtos.MinorType.INT) |
| .addArray("unnestColumn", TypeProtos.MinorType.VARCHAR).buildSchema(); |
| TupleMetadata[] incomingSchemas = {incomingSchema, incomingSchema}; |
| |
| Object[][][] baseline = { |
| { |
| {1, 1, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4, 4}, // rowNum |
| {"", "zero", "one", "two", "three", "four", "five", "six", "seven", "eight", "nine", "ten", "eleven", |
| "twelve"} // unnestColumn_flat |
| } |
| }; |
| |
| RecordBatch.IterOutcome[] iterOutcomes = {RecordBatch.IterOutcome.OK_NEW_SCHEMA, RecordBatch.IterOutcome.OK}; |
| |
| try { |
| testUnnest(incomingSchemas, iterOutcomes, data, baseline, false); |
| } catch (Exception e) { |
| fail("Failed due to exception: " + e.getMessage()); |
| } |
| } |
| |
| @Test |
| public void testUnnestMapColumn() { |
| |
| Object[][] data = getMapData(); |
| |
| // Create input schema |
| TupleMetadata incomingSchema = getRepeatedMapSchema(); |
| TupleMetadata[] incomingSchemas = {incomingSchema, incomingSchema}; |
| |
| Object[][][] baseline = getMapBaseline(); |
| |
| RecordBatch.IterOutcome[] iterOutcomes = {RecordBatch.IterOutcome.OK_NEW_SCHEMA, RecordBatch.IterOutcome.OK}; |
| |
| try { |
| testUnnest(incomingSchemas, iterOutcomes, data, baseline, false); |
| } catch (Exception e) { |
| fail("Failed due to exception: " + e.getMessage()); |
| } |
| } |
| |
| @Test |
| public void testUnnestEmptyList() { |
| |
| Object[][] data = { |
| { new String[] {}, |
| new String[] {} |
| }, |
| { new String[] {}, |
| new String[] {} |
| } |
| }; |
| |
| // Create input schema |
| TupleMetadata incomingSchema = new SchemaBuilder() |
| .add("someColumn", TypeProtos.MinorType.INT) |
| .addArray("unnestColumn", TypeProtos.MinorType.VARCHAR).buildSchema(); |
| TupleMetadata[] incomingSchemas = {incomingSchema, incomingSchema}; |
| |
| // All batches are empty |
| String[][][] baseline = {{{}}}; |
| |
| RecordBatch.IterOutcome[] iterOutcomes = {RecordBatch.IterOutcome.OK_NEW_SCHEMA, RecordBatch.IterOutcome.OK}; |
| |
| try { |
| testUnnest(incomingSchemas, iterOutcomes, data, baseline, false); |
| } catch (Exception e) { |
| fail("Failed due to exception: " + e.getMessage()); |
| } |
| } |
| |
| @Test |
| public void testUnnestMultipleNewSchemaIncoming() { |
| |
| // Schema changes in incoming have no effect on unnest unless the type of the |
| // unnest column itself has changed |
| Object[][] data = { |
| { |
| new String[] {"0", "1"}, |
| new String[] {"2", "3", "4"} |
| }, |
| { |
| new String[] {"5", "6" }, |
| }, |
| { |
| new String[] {"9"} |
| } |
| }; |
| |
| // Create input schema |
| TupleMetadata incomingSchema = new SchemaBuilder() |
| .add("someColumn", TypeProtos.MinorType.INT) |
| .addArray("unnestColumn", TypeProtos.MinorType.VARCHAR).buildSchema(); |
| TupleMetadata[] incomingSchemas = {incomingSchema, incomingSchema, incomingSchema}; |
| |
| Object[][][] baseline = { |
| { |
| {1, 1, 2, 2, 2, 3, 3}, |
| {"0", "1", "2", "3", "4", "5", "6"} |
| }, |
| { |
| {4}, |
| {"9"} |
| } |
| }; |
| |
| RecordBatch.IterOutcome[] iterOutcomes = { |
| RecordBatch.IterOutcome.OK_NEW_SCHEMA, |
| RecordBatch.IterOutcome.OK, |
| RecordBatch.IterOutcome.OK_NEW_SCHEMA}; |
| |
| try { |
| testUnnest(incomingSchemas, iterOutcomes, data, baseline, false); |
| } catch (Exception e) { |
| fail("Failed due to exception: " + e.getMessage()); |
| } |
| } |
| |
| @Test |
| public void testUnnestSchemaChange() { |
| Object[][] data = { |
| { |
| new String[] {"0", "1"}, |
| new String[] {"2", "3", "4"} |
| }, |
| { |
| new String[] {"5", "6" }, |
| }, |
| { |
| new int[] {9} |
| } |
| }; |
| |
| // Create input schema |
| TupleMetadata incomingSchema1 = new SchemaBuilder() |
| .add("someColumn", TypeProtos.MinorType.INT) |
| .addArray("unnestColumn", TypeProtos.MinorType.VARCHAR).buildSchema(); |
| TupleMetadata incomingSchema2 = new SchemaBuilder() |
| .add("someColumn", TypeProtos.MinorType.INT) |
| .addArray("unnestColumn", TypeProtos.MinorType.INT).buildSchema(); |
| |
| TupleMetadata[] incomingSchemas = {incomingSchema1, incomingSchema1, incomingSchema2}; |
| |
| Object[][][] baseline = { |
| { |
| {1, 1, 2, 2, 2, 3, 3}, |
| {"0", "1", "2", "3", "4", "5", "6"} |
| }, |
| { |
| {4}, |
| {9} |
| } |
| }; |
| |
| RecordBatch.IterOutcome[] iterOutcomes = { |
| RecordBatch.IterOutcome.OK_NEW_SCHEMA, |
| RecordBatch.IterOutcome.OK, |
| RecordBatch.IterOutcome.OK_NEW_SCHEMA}; |
| |
| try { |
| testUnnest(incomingSchemas, iterOutcomes, data, baseline, false); |
| } catch (Exception e) { |
| fail("Failed due to exception: " + e.getMessage()); |
| } |
| } |
| |
| private void testUnnestBatchSizing(int inputBatchSize, int limitOutputBatchSize, |
| int limitOutputBatchSizeBytes, boolean excludeUnnestColumn) { |
| // single record batch with single row. The unnest column has one |
| // more record than the batch size we want in the output |
| Object[][] data = new Object[1][1]; |
| |
| for (int i = 0; i < data.length; i++) { |
| for (int j = 0; j < data[i].length; j++) { |
| data[i][j] = new int[inputBatchSize]; |
| for (int k =0; k < inputBatchSize; k++) { |
| ((int[])data[i][j])[k] = k; |
| } |
| } |
| } |
| |
| Integer[][][] baseline = new Integer[2][2][]; |
| baseline[0][0] = new Integer[limitOutputBatchSize]; |
| baseline[0][1] = new Integer[limitOutputBatchSize]; |
| baseline[1][0] = new Integer[1]; |
| baseline[1][1] = new Integer[1]; |
| for (int i = 0; i < limitOutputBatchSize; i++) { |
| baseline[0][0][i] = 1; |
| baseline[0][1][i] = i; |
| } |
| baseline[1][0][0] = 1; // row Num |
| baseline[1][1][0] = limitOutputBatchSize; // value |
| |
| // Create input schema |
| TupleMetadata incomingSchema = new SchemaBuilder() |
| .add("rowNumber", TypeProtos.MinorType.INT) |
| .addArray("unnestColumn", TypeProtos.MinorType.INT).buildSchema(); |
| |
| TupleMetadata[] incomingSchemas = {incomingSchema}; |
| |
| RecordBatch.IterOutcome[] iterOutcomes = {RecordBatch.IterOutcome.OK}; |
| |
| final long outputBatchSize = fixture.getFragmentContext().getOptions().getOption(ExecConstants |
| .OUTPUT_BATCH_SIZE_VALIDATOR); |
| fixture.getFragmentContext().getOptions().setLocalOption(ExecConstants.OUTPUT_BATCH_SIZE, limitOutputBatchSizeBytes); |
| |
| try { |
| testUnnest(incomingSchemas, iterOutcomes, data, baseline, excludeUnnestColumn); |
| } catch (Exception e) { |
| fail("Failed due to exception: " + e.getMessage()); |
| } finally { |
| fixture.getFragmentContext().getOptions().setLocalOption(ExecConstants.OUTPUT_BATCH_SIZE, outputBatchSize); |
| } |
| } |
| |
| @Test |
| public void testUnnestLimitBatchSize_WithExcludedCols() { |
| LateralJoinPOP previoudPop = ljPopConfig; |
| List<SchemaPath> excludedCols = new ArrayList<>(); |
| excludedCols.add(SchemaPath.getSimplePath("unnestColumn")); |
| ljPopConfig = new LateralJoinPOP(null, null, JoinRelType.INNER, DrillLateralJoinRelBase.IMPLICIT_COLUMN, excludedCols); |
| final int limitedOutputBatchSize = 127; |
| final int inputBatchSize = limitedOutputBatchSize + 1; |
| // Since we want 127 row count and because of nearest power of 2 adjustment output row count will be reduced to |
| // 64. So we should configure batch size for (N+1) rows if we want to output N rows where N is not power of 2 |
| // size of lateral output batch = (N+1)*8 bytes, where N = output batch row count |
| // Lateral output batch size = (N+1) * (input row size without unnest field) + (N+1) * size of single unnest column |
| // = (N+1) * (size of row id) + (N+1) * (size of single array entry) |
| // = (N+1)*4 + (N+1) * 4 |
| // = (N+1) * 8 |
| // configure the output batch size to be one more record than that so that the batch sizer can round down |
| final int limitedOutputBatchSizeBytes = 8 * (limitedOutputBatchSize + 1); |
| testUnnestBatchSizing(inputBatchSize, limitedOutputBatchSize, limitedOutputBatchSizeBytes, true); |
| ljPopConfig = previoudPop; |
| } |
| |
| @Test |
| public void testUnnestLimitBatchSize() { |
| final int limitedOutputBatchSize = 127; |
| final int inputBatchSize = limitedOutputBatchSize + 1; |
| // size of lateral output batch = 4N * (N + 5) bytes, where N = output batch row count |
| // Lateral output batch size = N * input row size + N * size of single unnest column |
| // = N * (size of row id + size of array offset vector + (N + 1 )*size of single array entry)) |
| // + N * 4 |
| // = N * (4 + 2*4 + (N+1)*4 ) + N * 4 |
| // = N * (16 + 4N) + N * 4 |
| // = 4N * (N + 5) |
| // configure the output batch size to be one more record than that so that the batch sizer can round down |
| final int limitedOutputBatchSizeBytes = 4 * limitedOutputBatchSize * (limitedOutputBatchSize + 6); |
| testUnnestBatchSizing(inputBatchSize, limitedOutputBatchSize, limitedOutputBatchSizeBytes, false); |
| } |
| |
| @Test |
| // Limit sends a kill. Unnest has more than one record batch for a record when |
| // the kill is sent. |
| public void testUnnestKillFromLimitSubquery1() { |
| |
| // similar to previous test; we split a record across more than one batch. |
| // but we also set a limit less than the size of the batch so only one batch gets output. |
| final int limitedOutputBatchSize = 127; |
| final int inputBatchSize = limitedOutputBatchSize + 1; |
| final int limitedOutputBatchSizeBytes = 4 * limitedOutputBatchSize * (limitedOutputBatchSize + 6); |
| |
| // single record batch with single row. The unnest column has one |
| // more record than the batch size we want in the output |
| Object[][] data = new Object[1][1]; |
| |
| for (int i = 0; i < data.length; i++) { |
| for (int j = 0; j < data[i].length; j++) { |
| data[i][j] = new int[inputBatchSize]; |
| for (int k =0; k < inputBatchSize; k++) { |
| ((int[])data[i][j])[k] = k; |
| } |
| } |
| } |
| |
| // because of kill we only get one batch back |
| Integer[][][] baseline = new Integer[1][2][]; |
| baseline[0][0] = new Integer[limitedOutputBatchSize]; |
| baseline[0][1] = new Integer[limitedOutputBatchSize]; |
| for (int i = 0; i < limitedOutputBatchSize; i++) { |
| baseline[0][0][i] = 1; |
| baseline[0][1][i] = i; |
| } |
| |
| // Create input schema |
| TupleMetadata incomingSchema = new SchemaBuilder() |
| .add("rowNumber", TypeProtos.MinorType.INT) |
| .addArray("unnestColumn", TypeProtos.MinorType.INT).buildSchema(); |
| |
| TupleMetadata[] incomingSchemas = {incomingSchema}; |
| |
| RecordBatch.IterOutcome[] iterOutcomes = {RecordBatch.IterOutcome.OK}; |
| |
| final long outputBatchSize = fixture.getFragmentContext().getOptions().getOption(ExecConstants |
| .OUTPUT_BATCH_SIZE_VALIDATOR); |
| fixture.getFragmentContext().getOptions().setLocalOption(ExecConstants.OUTPUT_BATCH_SIZE, limitedOutputBatchSizeBytes); |
| |
| try { |
| testUnnest(incomingSchemas, iterOutcomes, -1, 1, data, baseline, false); // Limit of 100 values for unnest. |
| } catch (Exception e) { |
| fail("Failed due to exception: " + e.getMessage()); |
| } finally { |
| fixture.getFragmentContext().getOptions().setLocalOption(ExecConstants.OUTPUT_BATCH_SIZE, outputBatchSize); |
| } |
| } |
| |
| @Test |
| // Limit sends a kill. Unnest has exactly one record batch for a record when |
| // the kill is sent. This test is actually useless since it tests the behaviour of |
| // lateral which doesn't send kill at all if it gets an EMIT. We expect limit |
| // to do so, so let's keep the test to demonstrate the expected behaviour. |
| public void testUnnestKillFromLimitSubquery2() { |
| |
| // similar to previous test but the size of the array fits exactly into the record batch; |
| |
| final int limitedOutputBatchSize = 127; |
| final int inputBatchSize = limitedOutputBatchSize + 1; |
| final int limitedOutputBatchSizeBytes = 4 * limitedOutputBatchSize * (limitedOutputBatchSize + 6); |
| |
| // single record batch with single row. The unnest column has one |
| // more record than the batch size we want in the output |
| Object[][] data = new Object[1][1]; |
| |
| for (int i = 0; i < data.length; i++) { |
| for (int j = 0; j < data[i].length; j++) { |
| data[i][j] = new int[inputBatchSize]; |
| for (int k =0; k < inputBatchSize; k++) { |
| ((int[])data[i][j])[k] = k; |
| } |
| } |
| } |
| |
| // because of kill we only get one batch back |
| Integer[][][] baseline = new Integer[1][2][]; |
| baseline[0][0] = new Integer[limitedOutputBatchSize]; |
| baseline[0][1] = new Integer[limitedOutputBatchSize]; |
| for (int i = 0; i < limitedOutputBatchSize; i++) { |
| baseline[0][0][i] = 1; |
| baseline[0][1][i] = i; |
| } |
| |
| // Create input schema |
| TupleMetadata incomingSchema = new SchemaBuilder() |
| .add("rowNumber", TypeProtos.MinorType.INT) |
| .addArray("unnestColumn", TypeProtos.MinorType.INT).buildSchema(); |
| |
| TupleMetadata[] incomingSchemas = {incomingSchema}; |
| |
| RecordBatch.IterOutcome[] iterOutcomes = {RecordBatch.IterOutcome.OK}; |
| |
| final long outputBatchSize = fixture.getFragmentContext().getOptions().getOption(ExecConstants |
| .OUTPUT_BATCH_SIZE_VALIDATOR); |
| fixture.getFragmentContext().getOptions().setLocalOption(ExecConstants.OUTPUT_BATCH_SIZE, limitedOutputBatchSizeBytes); |
| |
| try { |
| testUnnest(incomingSchemas, iterOutcomes, -1, 1, data, baseline, false); // Limit of 100 values for unnest. |
| } catch (Exception e) { |
| fail("Failed due to exception: " + e.getMessage()); |
| } finally { |
| fixture.getFragmentContext().getOptions().setLocalOption(ExecConstants.OUTPUT_BATCH_SIZE, outputBatchSize); |
| } |
| } |
| |
| @Test |
| public void testUnnestNonArrayColumn() { |
| |
| Object[][] data = { |
| { new Integer (1), |
| new Integer (3)}, |
| { new Integer (6), |
| new Integer (10)} |
| }; |
| |
| // Create input schema |
| TupleMetadata incomingSchema = |
| new SchemaBuilder() |
| .add("rowNumber", TypeProtos.MinorType.INT) |
| .add("unnestColumn", TypeProtos.MinorType.INT) |
| .buildSchema(); |
| TupleMetadata[] incomingSchemas = { incomingSchema, incomingSchema }; |
| |
| // We expect an Exception |
| Integer[][][] baseline = {}; |
| |
| RecordBatch.IterOutcome[] iterOutcomes = {RecordBatch.IterOutcome.OK_NEW_SCHEMA, RecordBatch.IterOutcome.OK}; |
| |
| try { |
| testUnnest(incomingSchemas, iterOutcomes, data, baseline, false); |
| } catch (UserException|UnsupportedOperationException e) { |
| return; // succeeded |
| } catch (Exception e) { |
| fail("Failed due to exception: " + e.getMessage()); |
| } |
| } |
| |
| // test unnest for various input conditions without invoking kill |
| private <T> void testUnnest( |
| TupleMetadata[] incomingSchemas, |
| RecordBatch.IterOutcome[] iterOutcomes, |
| T[][] data, |
| T[][][] baseline, |
| boolean excludeUnnestColumn) throws Exception{ |
| testUnnest(incomingSchemas, iterOutcomes, -1, -1, data, baseline, excludeUnnestColumn); |
| } |
| |
| // test unnest for various input conditions optionally invoking kill. if the kill or killBatch |
| // parameter is greater than 0 then the record batch is sent a kill after that many batches have been processed |
| private <T> void testUnnest( TupleMetadata[] incomingSchemas, |
| RecordBatch.IterOutcome[] iterOutcomes, |
| int unnestLimit, // kill unnest after every 'unnestLimit' number of values in every record |
| int execKill, // number of batches after which to kill the execution (!) |
| T[][] data, |
| T[][][] baseline, |
| boolean excludeUnnestColumn) throws Exception { |
| |
| // Get the incoming container with dummy data for LJ |
| final List<VectorContainer> incomingContainer = new ArrayList<>(data.length); |
| |
| // Create data |
| ArrayList<RowSet.SingleRowSet> rowSets = new ArrayList<>(); |
| int rowNumber = 0; |
| int batchNum = 0; |
| for (Object[] recordBatch : data) { |
| RowSetBuilder rowSetBuilder = fixture.rowSetBuilder(incomingSchemas[batchNum]); |
| for ( Object rowData : recordBatch) { |
| rowSetBuilder.addRow(++rowNumber, rowData); |
| } |
| RowSet.SingleRowSet rowSet = rowSetBuilder.build(); |
| rowSets.add(rowSet); |
| incomingContainer.add(rowSet.container()); |
| batchNum++; |
| } |
| |
| // Get the unnest POPConfig |
| final UnnestPOP unnestPopConfig = new UnnestPOP(null, SchemaPath.getCompoundPath("unnestColumn"), DrillUnnestRelBase.IMPLICIT_COLUMN); |
| |
| // Get the IterOutcomes for LJ |
| final List<RecordBatch.IterOutcome> outcomes = new ArrayList<>(iterOutcomes.length); |
| for(RecordBatch.IterOutcome o : iterOutcomes) { |
| outcomes.add(o); |
| } |
| |
| // Create incoming MockRecordBatch |
| final MockRecordBatch incomingMockBatch = |
| new MockRecordBatch(fixture.getFragmentContext(), operatorContext, incomingContainer, outcomes, |
| incomingContainer.get(0).getSchema()); |
| |
| // setup Unnest record batch |
| final UnnestRecordBatch unnestBatch = |
| new UnnestRecordBatch(unnestPopConfig, fixture.getFragmentContext()); |
| |
| // project is required to rename the columns so as to disambiguate the same column name from |
| // unnest operator and the regular scan. |
| final Project projectPopConfig = new Project(DrillLogicalTestUtils.parseExprs("unnestColumn", "unnestColumn1", |
| unnestPopConfig.getImplicitColumn(), unnestPopConfig.getImplicitColumn()), null); |
| |
| final ProjectRecordBatch projectBatch = |
| new ProjectRecordBatch( projectPopConfig, unnestBatch, fixture.getFragmentContext()); |
| |
| final LateralJoinBatch lateralJoinBatch = |
| new LateralJoinBatch(ljPopConfig, fixture.getFragmentContext(), incomingMockBatch, projectBatch); |
| |
| // set pointer to Lateral in unnest |
| unnestBatch.setIncoming((LateralContract) lateralJoinBatch); |
| |
| // Simulate the pipeline by calling next on the incoming |
| |
| // results is an array of batches, each batch being an array of output vectors. |
| List<List<ValueVector> > resultList = new ArrayList<>(); |
| List<List<ValueVector> > results = null; |
| int batchesProcessed = 0; |
| try{ |
| try { |
| while (!isTerminal(lateralJoinBatch.next())) { |
| if (lateralJoinBatch.getRecordCount() > 0) { |
| addBatchToResults(resultList, lateralJoinBatch); |
| } |
| batchesProcessed++; |
| if (batchesProcessed == execKill) { |
| // Errors are reported by throwing an exception. |
| // Simulate by skipping out of the loop |
| break; |
| } |
| // else nothing to do |
| } |
| } catch (UserException e) { |
| throw e; |
| } catch (Exception e) { |
| fail(e.getMessage()); |
| } |
| |
| // Check results against baseline |
| results = resultList; |
| |
| int batchIndex = 0; |
| int vectorIndex = 0; |
| //int valueIndex = 0; |
| for (List<ValueVector> batch: results) { |
| int vectorCount= batch.size(); |
| int expectedVectorCount = (excludeUnnestColumn) ? 0 : 1; |
| expectedVectorCount += baseline[batchIndex].length; |
| if (vectorCount!= expectedVectorCount) { // baseline does not include the original unnest column |
| fail("Test failed in validating unnest output. Batch column count mismatch."); |
| } |
| for (ValueVector vv : batch) { |
| if(vv.getField().getName().equals("unnestColumn")) { |
| continue; // skip the original input column |
| } |
| int valueCount = vv.getAccessor().getValueCount(); |
| if (valueCount!= baseline[batchIndex][vectorIndex].length) { |
| fail("Test failed in validating unnest output. Value count mismatch in batch number " + (batchIndex+1) +"" |
| + "."); |
| } |
| |
| for (int valueIndex = 0; valueIndex < valueCount; valueIndex++) { |
| if (vv instanceof MapVector) { |
| if (!compareMapBaseline(baseline[batchIndex][vectorIndex][valueIndex], vv |
| .getAccessor() |
| .getObject(valueIndex))) { |
| fail("Test failed in validating unnest(Map) output. Value mismatch"); |
| } |
| } else if (vv instanceof VarCharVector) { |
| Object val = vv.getAccessor().getObject(valueIndex); |
| if (((String) baseline[batchIndex][vectorIndex][valueIndex]).compareTo(val.toString()) != 0) { |
| fail("Test failed in validating unnest output. Value mismatch. Baseline value[]" + vectorIndex + "][" + valueIndex |
| + "]" + ": " + baseline[vectorIndex][valueIndex] + " VV.getObject(valueIndex): " + val); |
| } |
| } else { |
| Object val = vv.getAccessor().getObject(valueIndex); |
| if (!baseline[batchIndex][vectorIndex][valueIndex].equals(val)) { |
| fail("Test failed in validating unnest output. Value mismatch. Baseline value[" + vectorIndex + "][" + valueIndex |
| + "]" + ": " |
| + baseline[batchIndex][vectorIndex][valueIndex] + " VV.getObject(valueIndex): " + val); |
| } |
| } |
| } |
| vectorIndex++; |
| } |
| vectorIndex=0; |
| batchIndex++; |
| } |
| } catch (UserException e) { |
| throw e; // Valid exception |
| } catch (Exception e) { |
| fail("Test failed. Exception : " + e.getMessage()); |
| } finally { |
| // Close all the resources for this test case |
| unnestBatch.close(); |
| lateralJoinBatch.close(); |
| incomingMockBatch.close(); |
| |
| if (results != null) { |
| for (List<ValueVector> batch : results) { |
| for (ValueVector vv : batch) { |
| vv.clear(); |
| } |
| } |
| } |
| for(RowSet.SingleRowSet rowSet: rowSets) { |
| rowSet.clear(); |
| } |
| } |
| } |
| |
| /** |
| * Build a schema with a repeated map - |
| * |
| * { |
| * rowNum, |
| * mapColumn : [ |
| * { |
| * colA, |
| * colB : [ |
| * varcharCol |
| * ] |
| * } |
| * ] |
| * } |
| * |
| * @see org.apache.drill.exec.physical.resultSet.impl.TestResultSetLoaderMapArray TestResultSetLoaderMapArray for |
| * similar schema and data |
| * @return TupleMetadata corresponding to the schema |
| */ |
| private TupleMetadata getRepeatedMapSchema() { |
| TupleMetadata schema = new SchemaBuilder() |
| .add("rowNum", TypeProtos.MinorType.INT) |
| .addMapArray("unnestColumn") |
| .add("colA", TypeProtos.MinorType.INT) |
| .addArray("colB", TypeProtos.MinorType.VARCHAR) |
| .resumeSchema() |
| .buildSchema(); |
| return schema; |
| } |
| |
| private Object[][] getMapData( ) { |
| |
| Object[][] d = { |
| { |
| new Object[] {}, |
| new Object[] { |
| new Object[] {11, new String[] {"1.1.1", "1.1.2" }}, |
| new Object[] {12, new String[] {"1.2.1", "1.2.2" }} |
| }, |
| |
| new Object[] { |
| new Object[] {21, new String[] {"2.1.1", "2.1.2" }}, |
| new Object[] {22, new String[] {}}, |
| new Object[] {23, new String[] {"2.3.1", "2.3.2" }} |
| } |
| }, |
| { |
| new Object[] { |
| new Object[] {31, new String[] {"3.1.1", "3.1.2" }}, |
| new Object[] {32, new String[] {"3.2.1", "3.2.2" }} |
| } |
| } |
| }; |
| |
| return d; |
| } |
| |
| private Object[][][] getMapBaseline() { |
| |
| Object[][][] d = { |
| { |
| {2,2,3,3,3,4,4}, |
| { |
| "{\"colA\":11,\"colB\":[\"1.1.1\",\"1.1.2\"]}", |
| "{\"colA\":12,\"colB\":[\"1.2.1\",\"1.2.2\"]}", |
| "{\"colA\":21,\"colB\":[\"2.1.1\",\"2.1.2\"]}", |
| "{\"colA\":22,\"colB\":[]}", |
| "{\"colA\":23,\"colB\":[\"2.3.1\",\"2.3.2\"]}", |
| "{\"colA\":31,\"colB\":[\"3.1.1\",\"3.1.2\"]}", |
| "{\"colA\":32,\"colB\":[\"3.2.1\",\"3.2.2\"]}" |
| } |
| } |
| }; |
| return d; |
| } |
| |
| private Object[][][] getNestedMapBaseline() { |
| |
| Object[][][] d = { |
| { |
| {2,2,2,2,3,3,3,3,4,4,4,4}, |
| { |
| "1.1.1", |
| "1.1.2", |
| "1.2.1", |
| "1.2.2", |
| "2.1.1", |
| "2.1.2", |
| "2.3.1", |
| "2.3.2", |
| "3.1.1", |
| "3.1.2", |
| "3.2.1", |
| "3.2.2" |
| } |
| } |
| }; |
| return d; |
| } |
| |
| private boolean compareMapBaseline(Object baselineValue, Object vector) { |
| String vv = vector.toString(); |
| String b = (String)baselineValue; |
| return vv.equalsIgnoreCase(b); |
| } |
| |
| private int addBatchToResults(List<List<ValueVector> > resultList, RecordBatch inputBatch) { |
| int count = 0; |
| final RecordBatchData batchCopy = new RecordBatchData(inputBatch, operatorContext.getAllocator()); |
| boolean success = false; |
| try { |
| count = batchCopy.getRecordCount(); |
| resultList.add(batchCopy.getVectors()); |
| success = true; |
| } finally { |
| if (!success) { |
| batchCopy.clear(); |
| } |
| } |
| return count; |
| } |
| |
| private boolean isTerminal(RecordBatch.IterOutcome outcome) { |
| return (outcome == RecordBatch.IterOutcome.NONE); |
| } |
| |
| /** |
| * Run a plan like the following for various input batches : |
| * Lateral1 |
| * / \ |
| * / Lateral2 |
| * Scan / \ |
| * / \ |
| * Project1 Project2 |
| * / \ |
| * / \ |
| * Unnest1 Unnest2 |
| * |
| * |
| * @param incomingSchemas |
| * @param iterOutcomes |
| * @param execKill |
| * @param data |
| * @param baseline |
| * @param <T> |
| * @throws Exception |
| */ |
| private <T> void testNestedUnnest( TupleMetadata[] incomingSchemas, |
| RecordBatch.IterOutcome[] iterOutcomes, |
| int execKill, // number of batches after which to kill the execution (!) |
| T[][] data, |
| T[][][] baseline) throws Exception { |
| |
| // Get the incoming container with dummy data for LJ |
| final List<VectorContainer> incomingContainer = new ArrayList<>(data.length); |
| |
| // Create data |
| ArrayList<RowSet.SingleRowSet> rowSets = new ArrayList<>(); |
| int rowNumber = 0; |
| int batchNum = 0; |
| for ( Object[] recordBatch : data) { |
| RowSetBuilder rowSetBuilder = fixture.rowSetBuilder(incomingSchemas[batchNum]); |
| for ( Object rowData : recordBatch) { |
| rowSetBuilder.addRow(++rowNumber, rowData); |
| } |
| RowSet.SingleRowSet rowSet = rowSetBuilder.build(); |
| rowSets.add(rowSet); |
| incomingContainer.add(rowSet.container()); |
| batchNum++; |
| } |
| |
| // Get the unnest POPConfig |
| final UnnestPOP unnestPopConfig1 = new UnnestPOP(null, SchemaPath.getSimplePath("unnestColumn"), DrillUnnestRelBase.IMPLICIT_COLUMN); |
| final UnnestPOP unnestPopConfig2 = new UnnestPOP(null, SchemaPath.getSimplePath("colB"), DrillUnnestRelBase.IMPLICIT_COLUMN); |
| |
| // Get the IterOutcomes for LJ |
| final List<RecordBatch.IterOutcome> outcomes = new ArrayList<>(iterOutcomes.length); |
| for(RecordBatch.IterOutcome o : iterOutcomes) { |
| outcomes.add(o); |
| } |
| |
| // Create incoming MockRecordBatch |
| final MockRecordBatch incomingMockBatch = |
| new MockRecordBatch(fixture.getFragmentContext(), operatorContext, incomingContainer, outcomes, |
| incomingContainer.get(0).getSchema()); |
| |
| // setup Unnest record batch |
| final UnnestRecordBatch unnestBatch1 = |
| new UnnestRecordBatch(unnestPopConfig1, fixture.getFragmentContext()); |
| final UnnestRecordBatch unnestBatch2 = |
| new UnnestRecordBatch(unnestPopConfig2, fixture.getFragmentContext()); |
| |
| // Create intermediate Project |
| final Project projectPopConfig1 = |
| new Project(DrillLogicalTestUtils.parseExprs("unnestColumn.colB", "colB", |
| unnestPopConfig1.getImplicitColumn(), unnestPopConfig1.getImplicitColumn()), unnestPopConfig1); |
| final ProjectRecordBatch projectBatch1 = |
| new ProjectRecordBatch(projectPopConfig1, unnestBatch1, fixture.getFragmentContext()); |
| final Project projectPopConfig2 = |
| new Project(DrillLogicalTestUtils.parseExprs("colB", "unnestColumn2", |
| unnestPopConfig2.getImplicitColumn(), unnestPopConfig2.getImplicitColumn()), unnestPopConfig2); |
| final ProjectRecordBatch projectBatch2 = |
| new ProjectRecordBatch(projectPopConfig2, unnestBatch2, fixture.getFragmentContext()); |
| |
| final LateralJoinPOP ljPopConfig2 = new LateralJoinPOP(projectPopConfig1, projectPopConfig2, JoinRelType.INNER, DrillLateralJoinRelBase.IMPLICIT_COLUMN, Lists.newArrayList()); |
| final LateralJoinPOP ljPopConfig1 = new LateralJoinPOP(mockPopConfig, ljPopConfig2, JoinRelType.INNER, DrillLateralJoinRelBase.IMPLICIT_COLUMN, Lists.newArrayList()); |
| |
| final LateralJoinBatch lateralJoinBatch2 = |
| new LateralJoinBatch(ljPopConfig2, fixture.getFragmentContext(), projectBatch1, projectBatch2); |
| final LateralJoinBatch lateralJoinBatch1 = |
| new LateralJoinBatch(ljPopConfig1, fixture.getFragmentContext(), incomingMockBatch, lateralJoinBatch2); |
| |
| // set pointer to Lateral in unnest |
| unnestBatch1.setIncoming((LateralContract) lateralJoinBatch1); |
| unnestBatch2.setIncoming((LateralContract) lateralJoinBatch2); |
| |
| // Simulate the pipeline by calling next on the incoming |
| |
| // results is an array ot batches, each batch being an array of output vectors. |
| List<List<ValueVector> > resultList = new ArrayList<>(); |
| List<List<ValueVector> > results = null; |
| int batchesProcessed = 0; |
| try{ |
| try { |
| while (!isTerminal(lateralJoinBatch1.next())) { |
| if (lateralJoinBatch1.getRecordCount() > 0) { |
| addBatchToResults(resultList, lateralJoinBatch1); |
| } |
| batchesProcessed++; |
| if (batchesProcessed == execKill) { |
| lateralJoinBatch1.getContext().getExecutorState().fail(new DrillException("Testing failure of execution.")); |
| lateralJoinBatch1.cancel(); |
| } |
| // else nothing to do |
| } |
| } catch (UserException e) { |
| throw e; |
| } catch (Exception e) { |
| throw new Exception ("Test failed to execute lateralJoinBatch.next() because: " + e.getMessage()); |
| } |
| |
| // Check results against baseline |
| results = resultList; |
| |
| int batchIndex = 0; |
| int vectorIndex = 0; |
| //int valueIndex = 0; |
| for ( List<ValueVector> batch: results) { |
| int vectorCount= batch.size(); |
| if (vectorCount!= baseline[batchIndex].length+2) { // baseline does not include the original unnest column(s) |
| fail("Test failed in validating unnest output. Batch column count mismatch."); |
| } |
| for (ValueVector vv : batch) { |
| if(vv.getField().getName().equals("unnestColumn") || vv.getField().getName().equals("colB")) { |
| continue; // skip the original input column |
| } |
| int valueCount = vv.getAccessor().getValueCount(); |
| if (valueCount!= baseline[batchIndex][vectorIndex].length) { |
| fail("Test failed in validating unnest output. Value count mismatch in batch number " + (batchIndex+1) +"" |
| + "."); |
| } |
| |
| for (int valueIndex = 0; valueIndex < valueCount; valueIndex++) { |
| if (vv instanceof MapVector) { |
| if (!compareMapBaseline(baseline[batchIndex][vectorIndex][valueIndex], vv |
| .getAccessor() |
| .getObject(valueIndex))) { |
| fail("Test failed in validating unnest(Map) output. Value mismatch"); |
| } |
| } else if (vv instanceof VarCharVector) { |
| Object val = vv.getAccessor().getObject(valueIndex); |
| if (((String) baseline[batchIndex][vectorIndex][valueIndex]).compareTo(val.toString()) != 0) { |
| fail("Test failed in validating unnest output. Value mismatch. Baseline value[]" + vectorIndex + "][" + valueIndex |
| + "]" + ": " + baseline[vectorIndex][valueIndex] + " VV.getObject(valueIndex): " + val); |
| } |
| } else { |
| Object val = vv.getAccessor().getObject(valueIndex); |
| if (!baseline[batchIndex][vectorIndex][valueIndex].equals(val)) { |
| fail("Test failed in validating unnest output. Value mismatch. Baseline value[" + vectorIndex + "][" + valueIndex |
| + "]" + ": " |
| + baseline[batchIndex][vectorIndex][valueIndex] + " VV.getObject(valueIndex): " + val); |
| } |
| } |
| } |
| vectorIndex++; |
| } |
| vectorIndex=0; |
| batchIndex++; |
| } |
| } catch (UserException e) { |
| throw e; // Valid exception |
| } catch (Exception e) { |
| fail("Test failed. Exception : " + e.getMessage()); |
| } finally { |
| // Close all the resources for this test case |
| unnestBatch1.close(); |
| lateralJoinBatch1.close(); |
| unnestBatch2.close(); |
| lateralJoinBatch2.close(); |
| incomingMockBatch.close(); |
| |
| if (results != null) { |
| for (List<ValueVector> batch : results) { |
| for (ValueVector vv : batch) { |
| vv.clear(); |
| } |
| } |
| } |
| for(RowSet.SingleRowSet rowSet: rowSets) { |
| rowSet.clear(); |
| } |
| } |
| } |
| |
| @Test |
| public void testNestedUnnestMapColumn() { |
| |
| Object[][] data = getMapData(); |
| |
| // Create input schema |
| TupleMetadata incomingSchema = getRepeatedMapSchema(); |
| TupleMetadata[] incomingSchemas = {incomingSchema, incomingSchema}; |
| |
| Object[][][] baseline = getNestedMapBaseline(); |
| |
| RecordBatch.IterOutcome[] iterOutcomes = {RecordBatch.IterOutcome.OK_NEW_SCHEMA, RecordBatch.IterOutcome.OK}; |
| |
| try { |
| testNestedUnnest(incomingSchemas, iterOutcomes, 0, data, baseline); |
| } catch (Exception e) { |
| fail("Failed due to exception: " + e.getMessage()); |
| } |
| } |
| } |
| |
