use of io.crate.analyze.AnalyzedCopyFromReturnSummary in project crate by crate.
the class CopyFromPlan method planCopyFromExecution.
public static ExecutionPlan planCopyFromExecution(AnalyzedCopyFrom copyFrom, DiscoveryNodes allNodes, PlannerContext context, Row params, SubQueryResults subQueryResults) {
var boundedCopyFrom = bind(copyFrom, context.transactionContext(), context.nodeContext(), params, subQueryResults);
/*
* Create a plan that reads json-objects-lines from a file
* and then executes upsert requests to index the data
*/
DocTableInfo table = boundedCopyFrom.tableInfo();
String partitionIdent = boundedCopyFrom.partitionIdent();
List<String> partitionedByNames = Collections.emptyList();
List<String> partitionValues = Collections.emptyList();
if (partitionIdent == null) {
if (table.isPartitioned()) {
partitionedByNames = Lists2.map(table.partitionedBy(), ColumnIdent::fqn);
}
} else {
assert table.isPartitioned() : "table must be partitioned if partitionIdent is set";
// partitionIdent is present -> possible to index raw source into concrete es index
partitionValues = PartitionName.decodeIdent(partitionIdent);
}
// need to exclude _id columns; they're auto generated and won't be available in the files being imported
ColumnIdent clusteredBy = table.clusteredBy();
if (DocSysColumns.ID.equals(clusteredBy)) {
clusteredBy = null;
}
List<Reference> primaryKeyRefs = table.primaryKey().stream().filter(r -> !r.equals(DocSysColumns.ID)).map(table::getReference).collect(Collectors.toList());
List<Symbol> toCollect = getSymbolsRequiredForShardIdCalc(primaryKeyRefs, table.partitionedByColumns(), clusteredBy == null ? null : table.getReference(clusteredBy));
Reference rawOrDoc = rawOrDoc(table, partitionIdent);
final int rawOrDocIdx = toCollect.size();
toCollect.add(rawOrDoc);
String[] excludes = partitionedByNames.size() > 0 ? partitionedByNames.toArray(new String[0]) : null;
InputColumns.SourceSymbols sourceSymbols = new InputColumns.SourceSymbols(toCollect);
Symbol clusteredByInputCol = null;
if (clusteredBy != null) {
clusteredByInputCol = InputColumns.create(table.getReference(clusteredBy), sourceSymbols);
}
SourceIndexWriterProjection sourceIndexWriterProjection;
List<? extends Symbol> projectionOutputs = AbstractIndexWriterProjection.OUTPUTS;
boolean returnSummary = copyFrom instanceof AnalyzedCopyFromReturnSummary;
boolean failFast = boundedCopyFrom.settings().getAsBoolean("fail_fast", false);
if (returnSummary || failFast) {
final InputColumn sourceUriSymbol = new InputColumn(toCollect.size(), DataTypes.STRING);
toCollect.add(SourceUriExpression.getReferenceForRelation(table.ident()));
final InputColumn sourceUriFailureSymbol = new InputColumn(toCollect.size(), DataTypes.STRING);
toCollect.add(SourceUriFailureExpression.getReferenceForRelation(table.ident()));
final InputColumn lineNumberSymbol = new InputColumn(toCollect.size(), DataTypes.LONG);
toCollect.add(SourceLineNumberExpression.getReferenceForRelation(table.ident()));
if (returnSummary) {
List<? extends Symbol> fields = ((AnalyzedCopyFromReturnSummary) copyFrom).outputs();
projectionOutputs = InputColumns.create(fields, new InputColumns.SourceSymbols(fields));
}
sourceIndexWriterProjection = new SourceIndexWriterReturnSummaryProjection(table.ident(), partitionIdent, table.getReference(DocSysColumns.RAW), new InputColumn(rawOrDocIdx, rawOrDoc.valueType()), table.primaryKey(), InputColumns.create(table.partitionedByColumns(), sourceSymbols), clusteredBy, boundedCopyFrom.settings(), null, excludes, InputColumns.create(primaryKeyRefs, sourceSymbols), clusteredByInputCol, projectionOutputs, // autoCreateIndices
table.isPartitioned(), sourceUriSymbol, sourceUriFailureSymbol, lineNumberSymbol);
} else {
sourceIndexWriterProjection = new SourceIndexWriterProjection(table.ident(), partitionIdent, table.getReference(DocSysColumns.RAW), new InputColumn(rawOrDocIdx, rawOrDoc.valueType()), table.primaryKey(), InputColumns.create(table.partitionedByColumns(), sourceSymbols), clusteredBy, boundedCopyFrom.settings(), null, excludes, InputColumns.create(primaryKeyRefs, sourceSymbols), clusteredByInputCol, projectionOutputs, // autoCreateIndices
table.isPartitioned());
}
// the partitionedBy-inputColumns created for the projection are still valid because the positions are not changed
if (partitionValues != null) {
rewriteToCollectToUsePartitionValues(table.partitionedByColumns(), partitionValues, toCollect);
}
FileUriCollectPhase collectPhase = new FileUriCollectPhase(context.jobId(), context.nextExecutionPhaseId(), "copyFrom", getExecutionNodes(allNodes, boundedCopyFrom.settings().getAsInt("num_readers", allNodes.getSize()), boundedCopyFrom.nodePredicate()), boundedCopyFrom.uri(), toCollect, Collections.emptyList(), boundedCopyFrom.settings().get("compression", null), boundedCopyFrom.settings().getAsBoolean("shared", null), CopyFromParserProperties.of(boundedCopyFrom.settings()), boundedCopyFrom.inputFormat(), boundedCopyFrom.settings());
Collect collect = new Collect(collectPhase, TopN.NO_LIMIT, 0, 1, -1, null);
// add the projection to the plan to ensure that the outputs are correctly set to the projection outputs
collect.addProjection(sourceIndexWriterProjection);
List<Projection> handlerProjections;
if (returnSummary) {
handlerProjections = Collections.emptyList();
} else {
handlerProjections = List.of(MergeCountProjection.INSTANCE);
}
return Merge.ensureOnHandler(collect, context, handlerProjections);
}
Aggregations