Search in sources :

Example 11 with VectorWrapper

use of org.apache.drill.exec.record.VectorWrapper in project drill by apache.

the class MergeJoinBatch method setRecordCountInContainer.

private void setRecordCountInContainer() {
    for (VectorWrapper vw : container) {
        Preconditions.checkArgument(!vw.isHyper());
        vw.getValueVector().getMutator().setValueCount(getRecordCount());
    }
}
Also used : VectorWrapper(org.apache.drill.exec.record.VectorWrapper)

Example 12 with VectorWrapper

use of org.apache.drill.exec.record.VectorWrapper in project drill by apache.

the class ProjectRecordBatch method classifyExpr.

private void classifyExpr(final NamedExpression ex, final RecordBatch incoming, final ClassifierResult result) {
    final NameSegment expr = ((SchemaPath) ex.getExpr()).getRootSegment();
    final NameSegment ref = ex.getRef().getRootSegment();
    final boolean exprHasPrefix = expr.getPath().contains(StarColumnHelper.PREFIX_DELIMITER);
    final boolean refHasPrefix = ref.getPath().contains(StarColumnHelper.PREFIX_DELIMITER);
    final boolean exprIsStar = expr.getPath().equals(StarColumnHelper.STAR_COLUMN);
    final boolean refContainsStar = ref.getPath().contains(StarColumnHelper.STAR_COLUMN);
    final boolean exprContainsStar = expr.getPath().contains(StarColumnHelper.STAR_COLUMN);
    final boolean refEndsWithStar = ref.getPath().endsWith(StarColumnHelper.STAR_COLUMN);
    String exprPrefix = EMPTY_STRING;
    String exprSuffix = expr.getPath();
    if (exprHasPrefix) {
        // get the prefix of the expr
        final String[] exprComponents = expr.getPath().split(StarColumnHelper.PREFIX_DELIMITER, 2);
        assert (exprComponents.length == 2);
        exprPrefix = exprComponents[0];
        exprSuffix = exprComponents[1];
        result.prefix = exprPrefix;
    }
    boolean exprIsFirstWildcard = false;
    if (exprContainsStar) {
        result.isStar = true;
        final Integer value = (Integer) result.prefixMap.get(exprPrefix);
        if (value == null) {
            final Integer n = 1;
            result.prefixMap.put(exprPrefix, n);
            exprIsFirstWildcard = true;
        } else {
            final Integer n = value + 1;
            result.prefixMap.put(exprPrefix, n);
        }
    }
    final int incomingSchemaSize = incoming.getSchema().getFieldCount();
    // input is '*' and output is 'prefix_*'
    if (exprIsStar && refHasPrefix && refEndsWithStar) {
        final String[] components = ref.getPath().split(StarColumnHelper.PREFIX_DELIMITER, 2);
        assert (components.length == 2);
        final String prefix = components[0];
        result.outputNames = Lists.newArrayList();
        for (final VectorWrapper<?> wrapper : incoming) {
            final ValueVector vvIn = wrapper.getValueVector();
            final String name = vvIn.getField().getPath();
            // add the prefix to the incoming column name
            final String newName = prefix + StarColumnHelper.PREFIX_DELIMITER + name;
            addToResultMaps(newName, result, false);
        }
    } else // input and output are the same
    if (expr.getPath().equalsIgnoreCase(ref.getPath()) && (!exprContainsStar || exprIsFirstWildcard)) {
        if (exprContainsStar && exprHasPrefix) {
            assert exprPrefix != null;
            int k = 0;
            result.outputNames = Lists.newArrayListWithCapacity(incomingSchemaSize);
            for (int j = 0; j < incomingSchemaSize; j++) {
                // initialize
                result.outputNames.add(EMPTY_STRING);
            }
            for (final VectorWrapper<?> wrapper : incoming) {
                final ValueVector vvIn = wrapper.getValueVector();
                final String incomingName = vvIn.getField().getPath();
                // get the prefix of the name
                final String[] nameComponents = incomingName.split(StarColumnHelper.PREFIX_DELIMITER, 2);
                // if incoming valuevector does not have a prefix, ignore it since this expression is not referencing it
                if (nameComponents.length <= 1) {
                    k++;
                    continue;
                }
                final String namePrefix = nameComponents[0];
                if (exprPrefix.equalsIgnoreCase(namePrefix)) {
                    final String newName = incomingName;
                    if (!result.outputMap.containsKey(newName)) {
                        result.outputNames.set(k, newName);
                        result.outputMap.put(newName, newName);
                    }
                }
                k++;
            }
        } else {
            result.outputNames = Lists.newArrayList();
            if (exprContainsStar) {
                for (final VectorWrapper<?> wrapper : incoming) {
                    final ValueVector vvIn = wrapper.getValueVector();
                    final String incomingName = vvIn.getField().getPath();
                    if (refContainsStar) {
                        // allow dups since this is likely top-level project
                        addToResultMaps(incomingName, result, true);
                    } else {
                        addToResultMaps(incomingName, result, false);
                    }
                }
            } else {
                final String newName = expr.getPath();
                if (!refHasPrefix && !exprHasPrefix) {
                    // allow dups since this is likely top-level project
                    addToResultMaps(newName, result, true);
                } else {
                    addToResultMaps(newName, result, false);
                }
            }
        }
    } else // input is wildcard and it is not the first wildcard
    if (exprIsStar) {
        result.outputNames = Lists.newArrayList();
        for (final VectorWrapper<?> wrapper : incoming) {
            final ValueVector vvIn = wrapper.getValueVector();
            final String incomingName = vvIn.getField().getPath();
            // allow dups since this is likely top-level project
            addToResultMaps(incomingName, result, true);
        }
    } else // only the output has prefix
    if (!exprHasPrefix && refHasPrefix) {
        result.outputNames = Lists.newArrayList();
        final String newName = ref.getPath();
        addToResultMaps(newName, result, false);
    } else // input has prefix but output does not
    if (exprHasPrefix && !refHasPrefix) {
        int k = 0;
        result.outputNames = Lists.newArrayListWithCapacity(incomingSchemaSize);
        for (int j = 0; j < incomingSchemaSize; j++) {
            // initialize
            result.outputNames.add(EMPTY_STRING);
        }
        for (final VectorWrapper<?> wrapper : incoming) {
            final ValueVector vvIn = wrapper.getValueVector();
            final String name = vvIn.getField().getPath();
            final String[] components = name.split(StarColumnHelper.PREFIX_DELIMITER, 2);
            if (components.length <= 1) {
                k++;
                continue;
            }
            final String namePrefix = components[0];
            final String nameSuffix = components[1];
            if (exprPrefix.equalsIgnoreCase(namePrefix)) {
                // // case insensitive matching of prefix.
                if (refContainsStar) {
                    // remove the prefix from the incoming column names
                    // for top level we need to make names unique
                    final String newName = getUniqueName(nameSuffix, result);
                    result.outputNames.set(k, newName);
                } else if (exprSuffix.equalsIgnoreCase(nameSuffix)) {
                    // case insensitive matching of field name.
                    // example: ref: $f1, expr: T0<PREFIX><column_name>
                    final String newName = ref.getPath();
                    result.outputNames.set(k, newName);
                }
            } else {
                result.outputNames.add(EMPTY_STRING);
            }
            k++;
        }
    } else // input and output have prefixes although they could be different...
    if (exprHasPrefix && refHasPrefix) {
        final String[] input = expr.getPath().split(StarColumnHelper.PREFIX_DELIMITER, 2);
        assert (input.length == 2);
        // not handled yet
        assert false : "Unexpected project expression or reference";
    } else {
        // if the incoming schema's column name matches the expression name of the Project,
        // then we just want to pick the ref name as the output column name
        result.outputNames = Lists.newArrayList();
        for (final VectorWrapper<?> wrapper : incoming) {
            final ValueVector vvIn = wrapper.getValueVector();
            final String incomingName = vvIn.getField().getPath();
            if (expr.getPath().equalsIgnoreCase(incomingName)) {
                // case insensitive matching of field name.
                final String newName = ref.getPath();
                addToResultMaps(newName, result, true);
            }
        }
    }
}
Also used : ValueVector(org.apache.drill.exec.vector.ValueVector) NameSegment(org.apache.drill.common.expression.PathSegment.NameSegment) SchemaPath(org.apache.drill.common.expression.SchemaPath) VectorWrapper(org.apache.drill.exec.record.VectorWrapper)

Example 13 with VectorWrapper

use of org.apache.drill.exec.record.VectorWrapper in project drill by apache.

the class ParquetRecordWriter method updateSchema.

@Override
public void updateSchema(VectorAccessible batch) throws IOException {
    if (this.batchSchema == null || !this.batchSchema.equals(batch.getSchema()) || containsComplexVectors(this.batchSchema)) {
        if (this.batchSchema != null) {
            flush();
        }
        this.batchSchema = batch.getSchema();
        newSchema();
    }
    TypedFieldId fieldId = batch.getValueVectorId(SchemaPath.getSimplePath(WriterPrel.PARTITION_COMPARATOR_FIELD));
    if (fieldId != null) {
        VectorWrapper w = batch.getValueAccessorById(BitVector.class, fieldId.getFieldIds());
        setPartitionVector((BitVector) w.getValueVector());
    }
}
Also used : TypedFieldId(org.apache.drill.exec.record.TypedFieldId) VectorWrapper(org.apache.drill.exec.record.VectorWrapper)

Example 14 with VectorWrapper

use of org.apache.drill.exec.record.VectorWrapper in project drill by apache.

the class ParquetResultListener method dataArrived.

@Override
public synchronized void dataArrived(QueryDataBatch result, ConnectionThrottle throttle) {
    logger.debug("result arrived in test batch listener.");
    int columnValCounter = 0;
    FieldInfo currentField;
    count += result.getHeader().getRowCount();
    boolean schemaChanged = false;
    final RecordBatchLoader batchLoader = new RecordBatchLoader(allocator);
    try {
        schemaChanged = batchLoader.load(result.getHeader().getDef(), result.getData());
    // TODO:  Clean:  DRILL-2933:  That load(...) no longer throws
    // SchemaChangeException, so check/clean catch clause below.
    } catch (SchemaChangeException e) {
        throw new RuntimeException(e);
    }
    // used to make sure each vector in the batch has the same number of records
    int valueCount = batchLoader.getRecordCount();
    // print headers.
    if (schemaChanged) {
    }
    for (final VectorWrapper vw : batchLoader) {
        final ValueVector vv = vw.getValueVector();
        currentField = props.fields.get(vv.getField().getPath());
        if (!valuesChecked.containsKey(vv.getField().getPath())) {
            valuesChecked.put(vv.getField().getPath(), 0);
            columnValCounter = 0;
        } else {
            columnValCounter = valuesChecked.get(vv.getField().getPath());
        }
        printColumnMajor(vv);
        if (testValues) {
            for (int j = 0; j < vv.getAccessor().getValueCount(); j++) {
                assertField(vv, j, currentField.type, currentField.values[columnValCounter % 3], currentField.name + "/");
                columnValCounter++;
            }
        } else {
            columnValCounter += vv.getAccessor().getValueCount();
        }
        valuesChecked.remove(vv.getField().getPath());
        assertEquals("Mismatched value count for vectors in the same batch.", valueCount, vv.getAccessor().getValueCount());
        valuesChecked.put(vv.getField().getPath(), columnValCounter);
    }
    if (ParquetRecordReaderTest.VERBOSE_DEBUG) {
        printRowMajor(batchLoader);
    }
    batchCounter++;
    batchLoader.clear();
    result.release();
}
Also used : ValueVector(org.apache.drill.exec.vector.ValueVector) SchemaChangeException(org.apache.drill.exec.exception.SchemaChangeException) RecordBatchLoader(org.apache.drill.exec.record.RecordBatchLoader) VectorWrapper(org.apache.drill.exec.record.VectorWrapper)

Example 15 with VectorWrapper

use of org.apache.drill.exec.record.VectorWrapper in project drill by apache.

the class TestOutputMutator method replace.

private void replace(ValueVector newVector, SchemaPath schemaPath) {
    List<ValueVector> vectors = Lists.newArrayList();
    for (VectorWrapper w : container) {
        ValueVector vector = w.getValueVector();
        if (vector.getField().getPath().equals(schemaPath)) {
            vectors.add(newVector);
        } else {
            vectors.add(w.getValueVector());
        }
        container.remove(vector);
    }
    container.addCollection(vectors);
}
Also used : ValueVector(org.apache.drill.exec.vector.ValueVector) VectorWrapper(org.apache.drill.exec.record.VectorWrapper)

Aggregations

VectorWrapper (org.apache.drill.exec.record.VectorWrapper)36 ValueVector (org.apache.drill.exec.vector.ValueVector)23 RecordBatchLoader (org.apache.drill.exec.record.RecordBatchLoader)17 Test (org.junit.Test)17 QueryDataBatch (org.apache.drill.exec.rpc.user.QueryDataBatch)16 DrillClient (org.apache.drill.exec.client.DrillClient)13 Drillbit (org.apache.drill.exec.server.Drillbit)13 RemoteServiceSet (org.apache.drill.exec.server.RemoteServiceSet)13 ExecTest (org.apache.drill.exec.ExecTest)4 SchemaChangeException (org.apache.drill.exec.exception.SchemaChangeException)4 Stopwatch (com.google.common.base.Stopwatch)3 IOException (java.io.IOException)3 SchemaPath (org.apache.drill.common.expression.SchemaPath)3 ClassTransformationException (org.apache.drill.exec.exception.ClassTransformationException)3 MaterializedField (org.apache.drill.exec.record.MaterializedField)3 TypedFieldId (org.apache.drill.exec.record.TypedFieldId)3 VectorContainer (org.apache.drill.exec.record.VectorContainer)3 VarBinaryHolder (org.apache.drill.exec.expr.holders.VarBinaryHolder)2 RecordBatchData (org.apache.drill.exec.physical.impl.sort.RecordBatchData)2 DrillbitEndpoint (org.apache.drill.exec.proto.CoordinationProtos.DrillbitEndpoint)2