Skip to content

Commit 04baff3

Browse files
pawel-big-lebowskipawel-big-lebowski
authored
include column lineage in dataset resource (#2148)
Signed-off-by: Pawel Leszczynski <leszczynski.pawel@gmail.com> Signed-off-by: Pawel Leszczynski <leszczynski.pawel@gmail.com>
1 parent e30a02b commit 04baff3

27 files changed

Lines changed: 426 additions & 157 deletions

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
* Implemented dataset symlink feature which allows providing multiple names for a dataset and adds edges to lineage graph based on symlinks [`#2066`](https://github.com/MarquezProject/marquez/pull/2066) [@pawel-big-lebowski](https://github.com/pawel-big-lebowski)
77
* Store column lineage facets in separate table [`#2096`](https://github.com/MarquezProject/marquez/pull/2096) [@mzareba382](https://github.com/mzareba382) [@pawel-big-lebowski](https://github.com/pawel-big-lebowski)
88
* Lineage graph endpoint for column lineage [`#2124`](https://github.com/MarquezProject/marquez/pull/2124) [@pawel-big-lebowski](https://github.com/pawel-big-lebowski)
9+
* Enrich returned dataset resource with column lineage information [`#2113`](https://github.com/MarquezProject/marquez/pull/2113) [@pawel-big-lebowski](https://github.com/pawel-big-lebowski)
910

1011
### Fixed
1112
* Add support for `parentRun` facet as reported by older Airflow OpenLineage versions [@collado-mike](https://github.com/collado-mike)

‎api/src/main/java/marquez/api/DatasetResource.java‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
import com.codahale.metrics.annotation.ResponseMetered;
1313
import com.codahale.metrics.annotation.Timed;
1414
import com.fasterxml.jackson.annotation.JsonProperty;
15+
import java.util.Arrays;
1516
import java.util.List;
1617
import java.util.Locale;
1718
import javax.validation.Valid;
@@ -85,10 +86,11 @@ public Response getDataset(
8586
@PathParam("dataset") DatasetName datasetName) {
8687
throwIfNotExists(namespaceName);
8788

88-
final Dataset dataset =
89+
Dataset dataset =
8990
datasetService
9091
.findWithTags(namespaceName.getValue(), datasetName.getValue())
9192
.orElseThrow(() -> new DatasetNotFoundException(datasetName));
93+
columnLineageService.enrichWithColumnLineage(Arrays.asList(dataset));
9294
return Response.ok(dataset).build();
9395
}
9496

@@ -147,6 +149,7 @@ public Response list(
147149

148150
final List<Dataset> datasets =
149151
datasetService.findAllWithTags(namespaceName.getValue(), limit, offset);
152+
columnLineageService.enrichWithColumnLineage(datasets);
150153
final int totalCount = datasetService.countFor(namespaceName.getValue());
151154
return Response.ok(new ResultsPage<>("datasets", datasets, totalCount)).build();
152155
}

‎api/src/main/java/marquez/db/ColumnLineageDao.java‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -153,4 +153,54 @@ Set<ColumnLineageNodeData> getLineage(
153153
int depth,
154154
@BindList(onEmpty = NULL_STRING) List<UUID> datasetFieldUuids,
155155
Instant createdAtUntil);
156+
157+
@SqlQuery(
158+
"""
159+
WITH selected_column_lineage AS (
160+
SELECT DISTINCT ON (cl.output_dataset_field_uuid, cl.input_dataset_field_uuid) cl.*
161+
FROM column_lineage cl
162+
JOIN dataset_fields df ON df.uuid = cl.output_dataset_field_uuid
163+
JOIN datasets_view dv ON dv.uuid = df.dataset_uuid
164+
WHERE ARRAY[<values>]::DATASET_NAME[] && dv.dataset_symlinks -- array of string pairs is cast onto array of DATASET_NAME types to be checked if it has non-empty intersection with dataset symlinks
165+
ORDER BY output_dataset_field_uuid, input_dataset_field_uuid, updated_at DESC, updated_at
166+
),
167+
dataset_fields_view AS (
168+
SELECT d.namespace_name as namespace_name, d.name as dataset_name, df.name as field_name, df.type, df.uuid
169+
FROM dataset_fields df
170+
INNER JOIN datasets_view d ON d.uuid = df.dataset_uuid
171+
)
172+
SELECT
173+
output_fields.namespace_name,
174+
output_fields.dataset_name,
175+
output_fields.field_name,
176+
output_fields.type,
177+
ARRAY_AGG(ARRAY[input_fields.namespace_name, input_fields.dataset_name, input_fields.field_name]) AS inputFields,
178+
c.transformation_description,
179+
c.transformation_type,
180+
c.created_at,
181+
c.updated_at
182+
FROM selected_column_lineage c
183+
INNER JOIN dataset_fields_view output_fields ON c.output_dataset_field_uuid = output_fields.uuid
184+
LEFT JOIN dataset_fields_view input_fields ON c.input_dataset_field_uuid = input_fields.uuid
185+
GROUP BY
186+
output_fields.namespace_name,
187+
output_fields.dataset_name,
188+
output_fields.field_name,
189+
output_fields.type,
190+
c.transformation_description,
191+
c.transformation_type,
192+
c.created_at,
193+
c.updated_at
194+
""")
195+
/**
196+
* Each dataset is identified by a pair of strings (namespace and name). A query returns column
197+
* lineage for multiple datasets, that's why a list of pairs is expected as an argument. "left"
198+
* and "right" properties correspond to Java Pair class properties defined to bind query template
199+
* with values
200+
*/
201+
Set<ColumnLineageNodeData> getLineageRowsForDatasets(
202+
@BindBeanList(
203+
propertyNames = {"left", "right"},
204+
value = "values")
205+
List<Pair<String, String>> datasets);
156206
}

‎api/src/main/java/marquez/db/DatasetFieldDao.java‎

Lines changed: 15 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -34,11 +34,14 @@
3434
@RegisterRowMapper(FieldDataMapper.class)
3535
public interface DatasetFieldDao extends BaseDao {
3636
@SqlQuery(
37-
"SELECT EXISTS ("
38-
+ "SELECT 1 FROM dataset_fields AS df "
39-
+ "INNER JOIN datasets_view AS d "
40-
+ " ON d.uuid = df.dataset_uuid AND d.name = :datasetName AND d.namespace_name = :namespaceName "
41-
+ "WHERE df.name = :name)")
37+
"""
38+
SELECT EXISTS (
39+
SELECT 1 FROM dataset_fields AS df
40+
INNER JOIN datasets_view AS d ON d.uuid = df.dataset_uuid
41+
WHERE CAST((:namespaceName, :datasetName) AS DATASET_NAME) = ANY(d.dataset_symlinks)
42+
AND df.name = :name
43+
)
44+
""")
4245
boolean exists(String namespaceName, String datasetName, String name);
4346

4447
default Dataset updateTags(
@@ -97,20 +100,20 @@ default Dataset updateTags(
97100
"""
98101
SELECT df.uuid
99102
FROM dataset_fields df
100-
INNER JOIN datasets_view AS d
101-
ON d.uuid = df.dataset_uuid AND d.name = :datasetName AND d.namespace_name = :namespace
103+
JOIN datasets_view AS d ON d.uuid = df.dataset_uuid
104+
WHERE CAST((:namespaceName, :datasetName) AS DATASET_NAME) = ANY(d.dataset_symlinks)
102105
""")
103-
List<UUID> findDatasetFieldsUuids(String namespace, String datasetName);
106+
List<UUID> findDatasetFieldsUuids(String namespaceName, String datasetName);
104107

105108
@SqlQuery(
106109
"""
107110
SELECT df.uuid
108111
FROM dataset_fields df
109-
INNER JOIN datasets_view AS d
110-
ON d.uuid = df.dataset_uuid AND d.name = :datasetName AND d.namespace_name = :namespace
111-
WHERE df.name = :name
112+
JOIN datasets_view AS d ON d.uuid = df.dataset_uuid
113+
WHERE CAST((:namespaceName, :datasetName) AS DATASET_NAME) = ANY(d.dataset_symlinks)
114+
AND df.name = :name
112115
""")
113-
Optional<UUID> findUuid(String namespace, String datasetName, String name);
116+
Optional<UUID> findUuid(String namespaceName, String datasetName, String name);
114117

115118
@SqlQuery(
116119
"SELECT f.*, "

‎api/src/main/java/marquez/service/ColumnLineageService.java‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,10 +22,14 @@
2222
import marquez.db.ColumnLineageDao;
2323
import marquez.db.DatasetFieldDao;
2424
import marquez.db.models.ColumnLineageNodeData;
25+
import marquez.service.models.ColumnLineage;
26+
import marquez.service.models.ColumnLineageInputField;
27+
import marquez.service.models.Dataset;
2528
import marquez.service.models.Edge;
2629
import marquez.service.models.Lineage;
2730
import marquez.service.models.Node;
2831
import marquez.service.models.NodeId;
32+
import org.apache.commons.lang3.tuple.Pair;
2933

3034
@Slf4j
3135
public class ColumnLineageService extends DelegatingDaos.DelegatingColumnLineageDao {
@@ -124,4 +128,51 @@ List<UUID> getColumnNodeUuids(NodeId nodeId) {
124128
}
125129
return columnNodeUuids;
126130
}
131+
132+
public void enrichWithColumnLineage(List<Dataset> datasets) {
133+
if (datasets.isEmpty()) {
134+
return;
135+
}
136+
137+
Set<ColumnLineageNodeData> lineageRowsForDatasets =
138+
getLineageRowsForDatasets(
139+
datasets.stream()
140+
.map(d -> Pair.of(d.getNamespace().getValue(), d.getName().getValue()))
141+
.collect(Collectors.toList()));
142+
143+
Map<Dataset, List<ColumnLineage>> datasetLineage = new HashMap<>();
144+
lineageRowsForDatasets.stream()
145+
.forEach(
146+
nodeData -> {
147+
Dataset dataset =
148+
datasets.stream()
149+
.filter(d -> d.getNamespace().getValue().equals(nodeData.getNamespace()))
150+
.filter(d -> d.getName().getValue().equals(nodeData.getDataset()))
151+
.findAny()
152+
.get();
153+
154+
if (!datasetLineage.containsKey(dataset)) {
155+
datasetLineage.put(dataset, new LinkedList<>());
156+
}
157+
datasetLineage
158+
.get(dataset)
159+
.add(
160+
ColumnLineage.builder()
161+
.name(nodeData.getField())
162+
.transformationDescription(nodeData.getTransformationDescription())
163+
.transformationType(nodeData.getTransformationType())
164+
.inputFields(
165+
nodeData.getInputFields().stream()
166+
.map(
167+
f ->
168+
new ColumnLineageInputField(
169+
f.getNamespace(), f.getDataset(), f.getField()))
170+
.collect(Collectors.toList()))
171+
.build());
172+
});
173+
174+
datasets.stream()
175+
.filter(dataset -> datasetLineage.containsKey(dataset))
176+
.forEach(dataset -> dataset.setColumnLineage(datasetLineage.get(dataset)));
177+
}
127178
}

‎api/src/main/java/marquez/service/models/ColumnLineage.java‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
/*
2+
* Copyright 2018-2022 contributors to the Marquez project
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package marquez.service.models;
7+
8+
import java.util.List;
9+
import javax.validation.constraints.NotNull;
10+
import lombok.Builder;
11+
import lombok.EqualsAndHashCode;
12+
import lombok.Getter;
13+
import lombok.ToString;
14+
15+
@EqualsAndHashCode
16+
@ToString
17+
@Builder
18+
@Getter
19+
public class ColumnLineage {
20+
@NotNull private String name;
21+
@NotNull private List<ColumnLineageInputField> inputFields;
22+
@NotNull private String transformationDescription;
23+
@NotNull private String transformationType;
24+
}

‎api/src/main/java/marquez/service/models/ColumnLineageInputField.java‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
/*
2+
* Copyright 2018-2022 contributors to the Marquez project
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package marquez.service.models;
7+
8+
import javax.validation.constraints.NotNull;
9+
import lombok.AllArgsConstructor;
10+
import lombok.EqualsAndHashCode;
11+
import lombok.Getter;
12+
import lombok.ToString;
13+
14+
@EqualsAndHashCode
15+
@ToString
16+
@Getter
17+
@AllArgsConstructor
18+
public class ColumnLineageInputField {
19+
@NotNull private String namespace;
20+
@NotNull private String dataset;
21+
@NotNull private String field;
22+
}

‎api/src/main/java/marquez/service/models/Dataset.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@ public abstract class Dataset {
5353
@Nullable private final String lastLifecycleState;
5454
@Nullable private final String description;
5555
@Nullable private final UUID currentVersion;
56+
@Getter @Setter @Nullable private List<ColumnLineage> columnLineage;
5657
@Getter ImmutableMap<String, Object> facets;
5758
@Getter private final boolean isDeleted;
5859

@@ -70,6 +71,7 @@ public Dataset(
7071
@Nullable final String lastLifecycleState,
7172
@Nullable final String description,
7273
@Nullable final UUID currentVersion,
74+
@Nullable final ImmutableList<ColumnLineage> columnLineage,
7375
@Nullable final ImmutableMap<String, Object> facets,
7476
boolean isDeleted) {
7577
this.id = id;
@@ -86,6 +88,7 @@ public Dataset(
8688
this.lastLifecycleState = lastLifecycleState;
8789
this.description = description;
8890
this.currentVersion = currentVersion;
91+
this.columnLineage = columnLineage;
8992
this.facets = (facets == null) ? ImmutableMap.of() : facets;
9093
this.isDeleted = isDeleted;
9194
}

‎api/src/main/java/marquez/service/models/DbTable.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@ public DbTable(
5353
lastLifecycleState,
5454
description,
5555
currentVersion,
56+
null,
5657
facets,
5758
isDeleted);
5859
}

‎api/src/main/java/marquez/service/models/Stream.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,7 @@ public Stream(
5959
lastLifecycleState,
6060
description,
6161
currentVersion,
62+
null,
6263
facets,
6364
isDeleted);
6465
this.schemaLocation = schemaLocation;

0 commit comments

Comments
 (0)
⚡