Examples with RowSetLoader org.apache.drill.exec.physical.resultSet.RowSetLoader used on opensource projects

Example 96 with RowSetLoader

use of org.apache.drill.exec.physical.resultSet.RowSetLoader in project drill by apache.

the class MetadataHandlerBatch method writeMetadataUsingBatchSchema.

private <T extends BaseMetadata & LocationProvider> VectorContainer writeMetadataUsingBatchSchema(List<T> metadataList) {
    Preconditions.checkArgument(!metadataList.isEmpty(), "Metadata list shouldn't be empty.");
    ResultSetLoader resultSetLoader = getResultSetLoaderWithBatchSchema();
    resultSetLoader.startBatch();
    RowSetLoader rowWriter = resultSetLoader.writer();
    Iterator<T> segmentsIterator = metadataList.iterator();
    while (!rowWriter.isFull() && segmentsIterator.hasNext()) {
        T metadata = segmentsIterator.next();
        metadataToHandle.remove(metadata.getMetadataInfo().identifier());
        List<Object> arguments = new ArrayList<>();
        for (VectorWrapper<?> vectorWrapper : container) {
            String[] identifierValues = Arrays.copyOf(MetadataIdentifierUtils.getValuesFromMetadataIdentifier(metadata.getMetadataInfo().identifier()), popConfig.getContext().segmentColumns().size());
            MaterializedField field = vectorWrapper.getField();
            String fieldName = field.getName();
            if (fieldName.equals(MetastoreAnalyzeConstants.LOCATION_FIELD)) {
                arguments.add(metadata.getPath().toUri().getPath());
            } else if (fieldName.equals(MetastoreAnalyzeConstants.LOCATIONS_FIELD)) {
                if (metadataType == MetadataType.SEGMENT) {
                    arguments.add(((SegmentMetadata) metadata).getLocations().stream().map(path -> path.toUri().getPath()).toArray(String[]::new));
                } else {
                    arguments.add(null);
                }
            } else if (popConfig.getContext().segmentColumns().contains(fieldName)) {
                arguments.add(identifierValues[popConfig.getContext().segmentColumns().indexOf(fieldName)]);
            } else if (AnalyzeColumnUtils.isColumnStatisticsField(fieldName)) {
                arguments.add(metadata.getColumnStatistics(SchemaPath.parseFromString(AnalyzeColumnUtils.getColumnName(fieldName))).get(AnalyzeColumnUtils.getStatisticsKind(fieldName)));
            } else if (AnalyzeColumnUtils.isMetadataStatisticsField(fieldName)) {
                arguments.add(metadata.getStatistic(AnalyzeColumnUtils.getStatisticsKind(fieldName)));
            } else if (fieldName.equals(MetastoreAnalyzeConstants.COLLECTED_MAP_FIELD)) {
                // collectedMap field value
                arguments.add(new Object[] {});
            } else if (fieldName.equals(MetastoreAnalyzeConstants.SCHEMA_FIELD)) {
                arguments.add(metadata.getSchema().jsonString());
            } else if (fieldName.equals(columnNamesOptions.lastModifiedTime())) {
                arguments.add(String.valueOf(metadata.getLastModifiedTime()));
            } else if (fieldName.equals(columnNamesOptions.rowGroupIndex())) {
                arguments.add(String.valueOf(((RowGroupMetadata) metadata).getRowGroupIndex()));
            } else if (fieldName.equals(columnNamesOptions.rowGroupStart())) {
                arguments.add(Long.toString(metadata.getStatistic(() -> ExactStatisticsConstants.START)));
            } else if (fieldName.equals(columnNamesOptions.rowGroupLength())) {
                arguments.add(Long.toString(metadata.getStatistic(() -> ExactStatisticsConstants.LENGTH)));
            } else if (fieldName.equals(MetastoreAnalyzeConstants.METADATA_TYPE)) {
                arguments.add(metadataType.name());
            } else {
                throw new UnsupportedOperationException(String.format("Found unexpected field [%s] in incoming batch.", field));
            }
        }
        rowWriter.addRow(arguments.toArray());
    }
    return resultSetLoader.harvest();
}

Example 97 with RowSetLoader

use of org.apache.drill.exec.physical.resultSet.RowSetLoader in project drill by apache.

the class JsonLoaderImpl method readBatch.

// JsonLoader
@Override
public boolean readBatch() {
    if (eof) {
        return false;
    }
    RowSetLoader rowWriter = rsLoader.writer();
    while (rowWriter.start()) {
        if (parser.next()) {
            // Add implicit fields
            if (implicitFields != null) {
                implicitFields.writeImplicitColumns();
            }
            rowWriter.save();
        } else if (rowWriter.limitReached(maxRows)) {
            eof = true;
            break;
        } else {
            // Special case for empty data sets to still record implicit columns
            if (implicitFields != null) {
                implicitFields.writeImplicitColumns();
                rowWriter.save();
            }
            eof = true;
            break;
        }
    }
    endBatch();
    return rsLoader.hasRows();
}

Example 98 with RowSetLoader

use of org.apache.drill.exec.physical.resultSet.RowSetLoader in project drill by apache.

the class ConstantColumnLoader method load.

@Override
public VectorContainer load(int rowCount) {
    loader.startBatch();
    RowSetLoader writer = loader.writer();
    for (int i = 0; i < rowCount; i++) {
        writer.start();
        loadRow(writer);
        writer.save();
    }
    return loader.harvest();
}