use of io.openlineage.client.OpenLineage.DatasetFacets in project OpenLineage by OpenLineage.
the class OpenLineageRunEventBuilder method buildInputDatasets.
private List<OpenLineage.InputDataset> buildInputDatasets(List<Object> nodes) {
openLineageContext.getQueryExecution().ifPresent(qe -> {
if (log.isDebugEnabled()) {
log.debug("Traversing optimized plan {}", qe.optimizedPlan().toJSON());
log.debug("Physical plan executed {}", qe.executedPlan().toJSON());
}
});
log.info("Visiting query plan {} with input dataset builders {}", openLineageContext.getQueryExecution(), inputDatasetBuilders);
Function1<LogicalPlan, Collection<InputDataset>> inputVisitor = visitLogicalPlan(PlanUtils.merge(inputDatasetQueryPlanVisitors));
List<OpenLineage.InputDataset> datasets = Stream.concat(buildDatasets(nodes, inputDatasetBuilders), openLineageContext.getQueryExecution().map(qe -> fromSeq(qe.optimizedPlan().map(inputVisitor)).stream().flatMap(Collection::stream).map(((Class<InputDataset>) InputDataset.class)::cast)).orElse(Stream.empty())).collect(Collectors.toList());
OpenLineage openLineage = openLineageContext.getOpenLineage();
if (!datasets.isEmpty()) {
Map<String, InputDatasetFacet> inputFacetsMap = new HashMap<>();
nodes.forEach(event -> inputDatasetFacetBuilders.forEach(fn -> fn.accept(event, inputFacetsMap::put)));
Map<String, DatasetFacets> datasetFacetsMap = new HashMap<>();
nodes.forEach(event -> inputDatasetFacetBuilders.forEach(fn -> fn.accept(event, inputFacetsMap::put)));
return datasets.stream().map(ds -> openLineage.newInputDatasetBuilder().name(ds.getName()).namespace(ds.getNamespace()).inputFacets(mergeFacets(inputFacetsMap, ds.getInputFacets(), InputDatasetInputFacets.class)).facets(mergeFacets(datasetFacetsMap, ds.getFacets(), DatasetFacets.class)).build()).collect(Collectors.toList());
}
return datasets;
}
Aggregations