use of org.apache.arrow.vector.ipc.ArrowWriter in project conquery by bakdata.
the class ResultArrowProcessor method getArrowResult.
public static <E extends ManagedExecution<?> & SingleTableResult> Response getArrowResult(Function<OutputStream, Function<VectorSchemaRoot, ArrowWriter>> writerProducer, Subject subject, E exec, Dataset dataset, DatasetRegistry datasetRegistry, boolean pretty, String fileExtension, MediaType mediaType, ConqueryConfig config) {
final Namespace namespace = datasetRegistry.get(dataset.getId());
ConqueryMDC.setLocation(subject.getName());
log.info("Downloading results for {} on dataset {}", exec, dataset);
subject.authorize(dataset, Ability.READ);
subject.authorize(dataset, Ability.DOWNLOAD);
subject.authorize(exec, Ability.READ);
// Check if subject is permitted to download on all datasets that were referenced by the query
authorizeDownloadDatasets(subject, exec);
if (!(exec instanceof ManagedQuery || (exec instanceof ManagedForm && ((ManagedForm) exec).getSubQueries().size() == 1))) {
return Response.status(HttpStatus.SC_UNPROCESSABLE_ENTITY, "Execution result is not a single Table").build();
}
// Get the locale extracted by the LocaleFilter
IdPrinter idPrinter = config.getFrontend().getQueryUpload().getIdPrinter(subject, exec, namespace);
final Locale locale = I18n.LOCALE.get();
PrintSettings settings = new PrintSettings(pretty, locale, datasetRegistry, config, idPrinter::createId);
// Collect ResultInfos for id columns and result columns
final List<ResultInfo> resultInfosId = config.getFrontend().getQueryUpload().getIdResultInfos();
final List<ResultInfo> resultInfosExec = exec.getResultInfos();
StreamingOutput out = output -> renderToStream(writerProducer.apply(output), settings, config.getArrow().getBatchSize(), resultInfosId, resultInfosExec, exec.streamResults());
return makeResponseWithFileName(out, exec.getLabelWithoutAutoLabelSuffix(), fileExtension, mediaType, ResultUtil.ContentDispositionOption.ATTACHMENT);
}
use of org.apache.arrow.vector.ipc.ArrowWriter in project conquery by bakdata.
the class ArrowRenderer method renderToStream.
public static void renderToStream(Function<VectorSchemaRoot, ArrowWriter> writerProducer, PrintSettings printSettings, int batchSize, List<ResultInfo> idHeaders, List<ResultInfo> resultInfo, Stream<EntityResult> results) throws IOException {
// Combine id and value Fields to one vector to build a schema
final UniqueNamer uniqNamer = new UniqueNamer(printSettings);
final List<Field> idFields = generateFields(idHeaders, uniqNamer);
List<Field> fields = new ArrayList<>(idFields);
fields.addAll(generateFields(resultInfo, uniqNamer));
VectorSchemaRoot root = VectorSchemaRoot.create(new Schema(fields, null), ROOT_ALLOCATOR);
// Build separate pipelines for id and value, as they have different sources but the same target
RowConsumer[] idWriters = generateWriterPipeline(root, 0, idHeaders.size(), printSettings, null);
RowConsumer[] valueWriter = generateWriterPipeline(root, idHeaders.size(), resultInfo.size(), printSettings, resultInfo);
// Write the data
try (ArrowWriter writer = writerProducer.apply(root)) {
write(writer, root, idWriters, valueWriter, printSettings.getIdMapper(), results, batchSize);
}
}
use of org.apache.arrow.vector.ipc.ArrowWriter in project snowflake-jdbc by snowflakedb.
the class SFArrowResultSetIT method createArrowFile.
private File createArrowFile(String fileName, Schema schema, Object[][] data, int rowsPerRecordBatch) throws IOException {
File file = resultFolder.newFile(fileName);
VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator);
try (ArrowWriter writer = new ArrowStreamWriter(root, new DictionaryProvider.MapDictionaryProvider(), new FileOutputStream(file))) {
writer.start();
for (int i = 0; i < data[0].length; ) {
int rowsToAppend = Math.min(rowsPerRecordBatch, data[0].length - i);
root.setRowCount(rowsToAppend);
for (int j = 0; j < data.length; j++) {
FieldVector vector = root.getFieldVectors().get(j);
switch(vector.getMinorType()) {
case INT:
writeIntToField(vector, data[j], i, rowsToAppend);
break;
}
}
writer.writeBatch();
i += rowsToAppend;
}
}
return file;
}
Aggregations