Examples with RecordBatchData org.apache.drill.exec.physical.impl.sort.RecordBatchData used on opensource projects

Example 6 with RecordBatchData

use of org.apache.drill.exec.physical.impl.sort.RecordBatchData in project drill by apache.

the class HashJoinBatch method executeBuildPhase.

public void executeBuildPhase() throws SchemaChangeException, ClassTransformationException, IOException {
    // skip first batch if count is zero, as it may be an empty schema batch
    if (right.getRecordCount() == 0) {
        for (final VectorWrapper<?> w : right) {
            w.clear();
        }
        rightUpstream = next(right);
    }
    boolean moreData = true;
    while (moreData) {
        switch(rightUpstream) {
            case OUT_OF_MEMORY:
            case NONE:
            case NOT_YET:
            case STOP:
                moreData = false;
                continue;
            case OK_NEW_SCHEMA:
                if (rightSchema == null) {
                    rightSchema = right.getSchema();
                    if (rightSchema.getSelectionVectorMode() != BatchSchema.SelectionVectorMode.NONE) {
                        final String errorMsg = new StringBuilder().append("Hash join does not support build batch with selection vectors. ").append("Build batch has selection mode = ").append(left.getSchema().getSelectionVectorMode()).toString();
                        throw new SchemaChangeException(errorMsg);
                    }
                    setupHashTable();
                } else {
                    if (!rightSchema.equals(right.getSchema())) {
                        throw SchemaChangeException.schemaChanged("Hash join does not support schema changes in build side.", rightSchema, right.getSchema());
                    }
                    hashTable.updateBatches();
                }
            // Fall through
            case OK:
                final int currentRecordCount = right.getRecordCount();
                /* For every new build batch, we store some state in the helper context
                     * Add new state to the helper context
                     */
                hjHelper.addNewBatch(currentRecordCount);
                // Holder contains the global index where the key is hashed into using the hash table
                final IndexPointer htIndex = new IndexPointer();
                // For every record in the build batch , hash the key columns
                for (int i = 0; i < currentRecordCount; i++) {
                    hashTable.put(i, htIndex, 1);
                    /* Use the global index returned by the hash table, to store
                         * the current record index and batch index. This will be used
                         * later when we probe and find a match.
                         */
                    hjHelper.setCurrentIndex(htIndex.value, buildBatchIndex, i);
                }
                /* Completed hashing all records in this batch. Transfer the batch
                     * to the hyper vector container. Will be used when we want to retrieve
                     * records that have matching keys on the probe side.
                     */
                final RecordBatchData nextBatch = new RecordBatchData(right, oContext.getAllocator());
                boolean success = false;
                try {
                    if (hyperContainer == null) {
                        hyperContainer = new ExpandableHyperContainer(nextBatch.getContainer());
                    } else {
                        hyperContainer.addBatch(nextBatch.getContainer());
                    }
                    // completed processing a batch, increment batch index
                    buildBatchIndex++;
                    success = true;
                } finally {
                    if (!success) {
                        nextBatch.clear();
                    }
                }
                break;
        }
        // Get the next record batch
        rightUpstream = next(HashJoinHelper.RIGHT_INPUT, right);
    }
}

Example 7 with RecordBatchData

use of org.apache.drill.exec.physical.impl.sort.RecordBatchData in project drill by apache.

the class ExternalSortBatch method innerNext.

@SuppressWarnings("resource")
@Override
public IterOutcome innerNext() {
    if (schema != null) {
        if (spillCount == 0) {
            return (getSelectionVector4().next()) ? IterOutcome.OK : IterOutcome.NONE;
        } else {
            Stopwatch w = Stopwatch.createStarted();
            int count = copier.next(targetRecordCount);
            if (count > 0) {
                long t = w.elapsed(TimeUnit.MICROSECONDS);
                logger.debug("Took {} us to merge {} records", t, count);
                container.setRecordCount(count);
                return IterOutcome.OK;
            } else {
                logger.debug("copier returned 0 records");
                return IterOutcome.NONE;
            }
        }
    }
    int totalCount = 0;
    // total number of batches received so far
    int totalBatches = 0;
    try {
        container.clear();
        outer: while (true) {
            IterOutcome upstream;
            if (first) {
                upstream = IterOutcome.OK_NEW_SCHEMA;
            } else {
                upstream = next(incoming);
            }
            if (upstream == IterOutcome.OK && sorter == null) {
                upstream = IterOutcome.OK_NEW_SCHEMA;
            }
            switch(upstream) {
                case NONE:
                    if (first) {
                        return upstream;
                    }
                    break outer;
                case NOT_YET:
                    throw new UnsupportedOperationException();
                case STOP:
                    return upstream;
                case OK_NEW_SCHEMA:
                case OK:
                    VectorContainer convertedBatch;
                    // only change in the case that the schema truly changes.  Artificial schema changes are ignored.
                    if (upstream == IterOutcome.OK_NEW_SCHEMA && !incoming.getSchema().equals(schema)) {
                        if (schema != null) {
                            if (unionTypeEnabled) {
                                this.schema = SchemaUtil.mergeSchemas(schema, incoming.getSchema());
                            } else {
                                throw SchemaChangeException.schemaChanged("Schema changes not supported in External Sort. Please enable Union type", schema, incoming.getSchema());
                            }
                        } else {
                            schema = incoming.getSchema();
                        }
                        convertedBatch = SchemaUtil.coerceContainer(incoming, schema, oContext);
                        for (BatchGroup b : batchGroups) {
                            b.setSchema(schema);
                        }
                        for (BatchGroup b : spilledBatchGroups) {
                            b.setSchema(schema);
                        }
                        this.sorter = createNewSorter(context, convertedBatch);
                    } else {
                        convertedBatch = SchemaUtil.coerceContainer(incoming, schema, oContext);
                    }
                    if (first) {
                        first = false;
                    }
                    if (convertedBatch.getRecordCount() == 0) {
                        for (VectorWrapper<?> w : convertedBatch) {
                            w.clear();
                        }
                        break;
                    }
                    SelectionVector2 sv2;
                    if (incoming.getSchema().getSelectionVectorMode() == BatchSchema.SelectionVectorMode.TWO_BYTE) {
                        sv2 = incoming.getSelectionVector2().clone();
                    } else {
                        try {
                            sv2 = newSV2();
                        } catch (InterruptedException e) {
                            return IterOutcome.STOP;
                        } catch (OutOfMemoryException e) {
                            throw new OutOfMemoryException(e);
                        }
                    }
                    int count = sv2.getCount();
                    totalCount += count;
                    totalBatches++;
                    sorter.setup(context, sv2, convertedBatch);
                    sorter.sort(sv2);
                    RecordBatchData rbd = new RecordBatchData(convertedBatch, oAllocator);
                    boolean success = false;
                    try {
                        rbd.setSv2(sv2);
                        batchGroups.add(new BatchGroup(rbd.getContainer(), rbd.getSv2(), oContext));
                        if (peakNumBatches < batchGroups.size()) {
                            peakNumBatches = batchGroups.size();
                            stats.setLongStat(Metric.PEAK_BATCHES_IN_MEMORY, peakNumBatches);
                        }
                        batchesSinceLastSpill++;
                        if (// If we haven't spilled so far, do we have enough memory for MSorter if this turns out to be the last incoming batch?
                        (spillCount == 0 && !hasMemoryForInMemorySort(totalCount)) || // If we haven't spilled so far, make sure we don't exceed the maximum number of batches SV4 can address
                        (spillCount == 0 && totalBatches > Character.MAX_VALUE) || // current memory used is more than 95% of memory usage limit of this operator
                        (oAllocator.getAllocatedMemory() > .95 * oAllocator.getLimit()) || // since the last spill exceed the defined limit
                        (batchGroups.size() > SPILL_THRESHOLD && batchesSinceLastSpill >= SPILL_BATCH_GROUP_SIZE)) {
                            if (firstSpillBatchCount == 0) {
                                firstSpillBatchCount = batchGroups.size();
                            }
                            if (spilledBatchGroups.size() > firstSpillBatchCount / 2) {
                                logger.info("Merging spills");
                                final BatchGroup merged = mergeAndSpill(spilledBatchGroups);
                                if (merged != null) {
                                    spilledBatchGroups.addFirst(merged);
                                }
                            }
                            final BatchGroup merged = mergeAndSpill(batchGroups);
                            if (merged != null) {
                                // make sure we don't add null to spilledBatchGroups
                                spilledBatchGroups.add(merged);
                                batchesSinceLastSpill = 0;
                            }
                        }
                        success = true;
                    } finally {
                        if (!success) {
                            rbd.clear();
                        }
                    }
                    break;
                case OUT_OF_MEMORY:
                    logger.debug("received OUT_OF_MEMORY, trying to spill");
                    if (batchesSinceLastSpill > 2) {
                        final BatchGroup merged = mergeAndSpill(batchGroups);
                        if (merged != null) {
                            spilledBatchGroups.add(merged);
                            batchesSinceLastSpill = 0;
                        }
                    } else {
                        logger.debug("not enough batches to spill, sending OUT_OF_MEMORY downstream");
                        return IterOutcome.OUT_OF_MEMORY;
                    }
                    break;
                default:
                    throw new UnsupportedOperationException();
            }
        }
        if (totalCount == 0) {
            return IterOutcome.NONE;
        }
        if (spillCount == 0) {
            if (builder != null) {
                builder.clear();
                builder.close();
            }
            builder = new SortRecordBatchBuilder(oAllocator);
            for (BatchGroup group : batchGroups) {
                RecordBatchData rbd = new RecordBatchData(group.getContainer(), oAllocator);
                rbd.setSv2(group.getSv2());
                builder.add(rbd);
            }
            builder.build(context, container);
            sv4 = builder.getSv4();
            mSorter = createNewMSorter();
            mSorter.setup(context, oAllocator, getSelectionVector4(), this.container);
            // For testing memory-leak purpose, inject exception after mSorter finishes setup
            injector.injectUnchecked(context.getExecutionControls(), INTERRUPTION_AFTER_SETUP);
            mSorter.sort(this.container);
            // sort may have prematurely exited due to should continue returning false.
            if (!context.shouldContinue()) {
                return IterOutcome.STOP;
            }
            // For testing memory-leak purpose, inject exception after mSorter finishes sorting
            injector.injectUnchecked(context.getExecutionControls(), INTERRUPTION_AFTER_SORT);
            sv4 = mSorter.getSV4();
            container.buildSchema(SelectionVectorMode.FOUR_BYTE);
        } else {
            // some batches were spilled
            final BatchGroup merged = mergeAndSpill(batchGroups);
            if (merged != null) {
                spilledBatchGroups.add(merged);
            }
            batchGroups.addAll(spilledBatchGroups);
            // no need to cleanup spilledBatchGroups, all it's batches are in batchGroups now
            spilledBatchGroups = null;
            logger.warn("Starting to merge. {} batch groups. Current allocated memory: {}", batchGroups.size(), oAllocator.getAllocatedMemory());
            VectorContainer hyperBatch = constructHyperBatch(batchGroups);
            createCopier(hyperBatch, batchGroups, container, false);
            int estimatedRecordSize = 0;
            for (VectorWrapper<?> w : batchGroups.get(0)) {
                try {
                    estimatedRecordSize += TypeHelper.getSize(w.getField().getType());
                } catch (UnsupportedOperationException e) {
                    estimatedRecordSize += 50;
                }
            }
            targetRecordCount = Math.min(MAX_BATCH_SIZE, Math.max(1, COPIER_BATCH_MEM_LIMIT / estimatedRecordSize));
            int count = copier.next(targetRecordCount);
            container.buildSchema(SelectionVectorMode.NONE);
            container.setRecordCount(count);
        }
        return IterOutcome.OK_NEW_SCHEMA;
    } catch (SchemaChangeException ex) {
        kill(false);
        context.fail(UserException.unsupportedError(ex).message("Sort doesn't currently support sorts with changing schemas").build(logger));
        return IterOutcome.STOP;
    } catch (ClassTransformationException | IOException ex) {
        kill(false);
        context.fail(ex);
        return IterOutcome.STOP;
    } catch (UnsupportedOperationException e) {
        throw new RuntimeException(e);
    }
}

Example 8 with RecordBatchData

use of org.apache.drill.exec.physical.impl.sort.RecordBatchData in project drill by apache.

the class RecordIterator method forward.

// Move forward by delta (may cross one or more record batches)
public void forward(long delta) {
    if (!enableMarkAndReset) {
        throw new UnsupportedOperationException("mark and reset disabled for this RecordIterator");
    }
    assert delta >= 0;
    assert (delta + outerPosition) < totalRecordCount;
    final long nextOuterPosition = delta + outerPosition;
    final RecordBatchData rbdNew = batches.get(nextOuterPosition);
    final RecordBatchData rbdOld = batches.get(outerPosition);
    assert rbdNew != null;
    assert rbdOld != null;
    container.transferOut(rbdOld.getContainer());
    // Get vectors from new position.
    container.transferIn(rbdNew.getContainer());
    outerPosition = nextOuterPosition;
    final Range<Long> markedBatchRange = batches.getEntry(outerPosition).getKey();
    startBatchPosition = markedBatchRange.lowerEndpoint();
    innerPosition = (int) (outerPosition - startBatchPosition);
    innerRecordCount = (int) (markedBatchRange.upperEndpoint() - startBatchPosition);
}

Example 9 with RecordBatchData

use of org.apache.drill.exec.physical.impl.sort.RecordBatchData in project drill by apache.

the class RecordIterator method next.

/**
   * Move iterator to next record.
   * @return
   *     Status of current record batch read.
   */
public IterOutcome next() {
    if (finished()) {
        return lastOutcome;
    }
    long nextOuterPosition = outerPosition + 1;
    final int nextInnerPosition = innerPosition + 1;
    if (!initialized || nextOuterPosition >= totalRecordCount) {
        nextBatch();
        switch(lastOutcome) {
            case NONE:
            case STOP:
                // No more data, disallow reads unless reset is called.
                outerPosition = nextOuterPosition;
                lastBatchRead = true;
                if (!enableMarkAndReset) {
                    container.clear();
                }
                break;
            case OK_NEW_SCHEMA:
            case OK:
                // If Schema changes in the middle of the execution clear out data.
                if (initialized && lastOutcome == IterOutcome.OK_NEW_SCHEMA) {
                    clear();
                    resetIndices();
                    initialized = false;
                    nextOuterPosition = 0;
                }
                // Transfer vectors from incoming record batch.
                final RecordBatchData rbd = new RecordBatchData(incoming, oContext.getAllocator());
                innerRecordCount = incoming.getRecordCount();
                if (!initialized) {
                    for (VectorWrapper<?> w : rbd.getContainer()) {
                        container.addOrGet(w.getField());
                    }
                    container.buildSchema(rbd.getContainer().getSchema().getSelectionVectorMode());
                    initialized = true;
                }
                if (innerRecordCount > 0) {
                    if (enableMarkAndReset) {
                        // Transfer vectors back to old batch.
                        if (startBatchPosition != -1 && batches.get(startBatchPosition) != null) {
                            container.transferOut(batches.get(outerPosition).getContainer());
                        }
                        container.transferIn(rbd.getContainer());
                        batches.put(Range.closedOpen(nextOuterPosition, nextOuterPosition + innerRecordCount), rbd);
                    } else {
                        container.zeroVectors();
                        container.transferIn(rbd.getContainer());
                    }
                    innerPosition = 0;
                    startBatchPosition = nextOuterPosition;
                    outerPosition = nextOuterPosition;
                    totalRecordCount += innerRecordCount;
                } else {
                    // Release schema/empty batches.
                    rbd.clear();
                }
                break;
            case OUT_OF_MEMORY:
                return lastOutcome;
            case NOT_YET:
            default:
                throw new UnsupportedOperationException("Unsupported outcome received " + lastOutcome);
        }
    } else {
        if (nextInnerPosition >= innerRecordCount) {
            assert enableMarkAndReset;
            // move to next batch
            final RecordBatchData rbdNew = batches.get(nextOuterPosition);
            final RecordBatchData rbdOld = batches.get(outerPosition);
            assert rbdNew != null;
            assert rbdOld != null;
            assert rbdOld != rbdNew;
            container.transferOut(rbdOld.getContainer());
            container.transferIn(rbdNew.getContainer());
            innerPosition = 0;
            outerPosition = nextOuterPosition;
            startBatchPosition = batches.getEntry(outerPosition).getKey().lowerEndpoint();
            innerRecordCount = (int) (batches.getEntry(outerPosition).getKey().upperEndpoint() - startBatchPosition);
        } else {
            outerPosition = nextOuterPosition;
            innerPosition = nextInnerPosition;
        }
    }
    return lastOutcome;
}

Example 10 with RecordBatchData

use of org.apache.drill.exec.physical.impl.sort.RecordBatchData in project drill by apache.

the class RecordIterator method mark.

public void mark() {
    if (!enableMarkAndReset) {
        throw new UnsupportedOperationException("mark and reset disabled for this RecordIterator");
    }
    // Release all batches before current batch. [0 to startBatchPosition).
    final Map<Range<Long>, RecordBatchData> oldBatches = batches.subRangeMap(Range.closedOpen(0l, startBatchPosition)).asMapOfRanges();
    for (RecordBatchData rbd : oldBatches.values()) {
        rbd.clear();
    }
    batches.remove(Range.closedOpen(0l, startBatchPosition));
    markedInnerPosition = innerPosition;
    markedOuterPosition = outerPosition;
}