Skip to content

Commit b7bb5ca

Browse files
asikowitzasikowitz
authored
feat(graphql/ml): Add custom properties to ml entities (#12152)
1 parent 9031b49 commit b7bb5ca

11 files changed

+108
-52
lines changed

‎datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mappers/EmbeddedModelMapper.java

+12
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
package com.linkedin.datahub.graphql.types.mappers;
2+
3+
import com.linkedin.common.urn.Urn;
4+
import com.linkedin.datahub.graphql.QueryContext;
5+
import javax.annotation.Nonnull;
6+
import javax.annotation.Nullable;
7+
8+
/** Made for models that are embedded in other models and thus do not encode their own URN. */
9+
public interface EmbeddedModelMapper<I, O> {
10+
O apply(
11+
@Nullable final QueryContext context, @Nonnull final I input, @Nonnull final Urn entityUrn);
12+
}

‎datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLFeatureMapper.java

+8-4
Original file line numberDiff line numberDiff line change
@@ -75,7 +75,8 @@ public MLFeature apply(
7575
mlFeature.setOwnership(
7676
OwnershipMapper.map(context, new Ownership(dataMap), entityUrn)));
7777
mappingHelper.mapToResult(
78-
context, ML_FEATURE_PROPERTIES_ASPECT_NAME, MLFeatureMapper::mapMLFeatureProperties);
78+
ML_FEATURE_PROPERTIES_ASPECT_NAME,
79+
(entity, dataMap) -> mapMLFeatureProperties(context, entity, dataMap, entityUrn));
7980
mappingHelper.mapToResult(
8081
INSTITUTIONAL_MEMORY_ASPECT_NAME,
8182
(mlFeature, dataMap) ->
@@ -138,10 +139,13 @@ private static void mapMLFeatureKey(@Nonnull MLFeature mlFeature, @Nonnull DataM
138139
private static void mapMLFeatureProperties(
139140
@Nullable final QueryContext context,
140141
@Nonnull MLFeature mlFeature,
141-
@Nonnull DataMap dataMap) {
142+
@Nonnull DataMap dataMap,
143+
@Nonnull Urn entityUrn) {
142144
MLFeatureProperties featureProperties = new MLFeatureProperties(dataMap);
143-
mlFeature.setFeatureProperties(MLFeaturePropertiesMapper.map(context, featureProperties));
144-
mlFeature.setProperties(MLFeaturePropertiesMapper.map(context, featureProperties));
145+
com.linkedin.datahub.graphql.generated.MLFeatureProperties graphqlProperties =
146+
MLFeaturePropertiesMapper.map(context, featureProperties, entityUrn);
147+
mlFeature.setFeatureProperties(graphqlProperties);
148+
mlFeature.setProperties(graphqlProperties);
145149
mlFeature.setDescription(featureProperties.getDescription());
146150
if (featureProperties.getDataType() != null) {
147151
mlFeature.setDataType(MLFeatureDataType.valueOf(featureProperties.getDataType().toString()));

‎datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLFeaturePropertiesMapper.java

+14-6
Original file line numberDiff line numberDiff line change
@@ -1,29 +1,34 @@
11
package com.linkedin.datahub.graphql.types.mlmodel.mappers;
22

3+
import com.linkedin.common.urn.Urn;
34
import com.linkedin.datahub.graphql.QueryContext;
45
import com.linkedin.datahub.graphql.generated.Dataset;
56
import com.linkedin.datahub.graphql.generated.MLFeatureDataType;
67
import com.linkedin.datahub.graphql.generated.MLFeatureProperties;
7-
import com.linkedin.datahub.graphql.types.mappers.ModelMapper;
8+
import com.linkedin.datahub.graphql.types.common.mappers.CustomPropertiesMapper;
9+
import com.linkedin.datahub.graphql.types.mappers.EmbeddedModelMapper;
810
import java.util.stream.Collectors;
11+
import javax.annotation.Nonnull;
912
import javax.annotation.Nullable;
10-
import lombok.NonNull;
1113

1214
public class MLFeaturePropertiesMapper
13-
implements ModelMapper<com.linkedin.ml.metadata.MLFeatureProperties, MLFeatureProperties> {
15+
implements EmbeddedModelMapper<
16+
com.linkedin.ml.metadata.MLFeatureProperties, MLFeatureProperties> {
1417

1518
public static final MLFeaturePropertiesMapper INSTANCE = new MLFeaturePropertiesMapper();
1619

1720
public static MLFeatureProperties map(
1821
@Nullable QueryContext context,
19-
@NonNull final com.linkedin.ml.metadata.MLFeatureProperties mlFeatureProperties) {
20-
return INSTANCE.apply(context, mlFeatureProperties);
22+
@Nonnull final com.linkedin.ml.metadata.MLFeatureProperties mlFeatureProperties,
23+
@Nonnull Urn entityUrn) {
24+
return INSTANCE.apply(context, mlFeatureProperties, entityUrn);
2125
}
2226

2327
@Override
2428
public MLFeatureProperties apply(
2529
@Nullable QueryContext context,
26-
@NonNull final com.linkedin.ml.metadata.MLFeatureProperties mlFeatureProperties) {
30+
@Nonnull final com.linkedin.ml.metadata.MLFeatureProperties mlFeatureProperties,
31+
@Nonnull Urn entityUrn) {
2732
final MLFeatureProperties result = new MLFeatureProperties();
2833

2934
result.setDescription(mlFeatureProperties.getDescription());
@@ -45,6 +50,9 @@ public MLFeatureProperties apply(
4550
.collect(Collectors.toList()));
4651
}
4752

53+
result.setCustomProperties(
54+
CustomPropertiesMapper.map(mlFeatureProperties.getCustomProperties(), entityUrn));
55+
4856
return result;
4957
}
5058
}

‎datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLFeatureTableMapper.java

+5-5
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,7 @@ public MLFeatureTable apply(
7676
mappingHelper.mapToResult(ML_FEATURE_TABLE_KEY_ASPECT_NAME, this::mapMLFeatureTableKey);
7777
mappingHelper.mapToResult(
7878
ML_FEATURE_TABLE_PROPERTIES_ASPECT_NAME,
79-
(entity, dataMap) -> this.mapMLFeatureTableProperties(context, entity, dataMap, entityUrn));
79+
(entity, dataMap) -> mapMLFeatureTableProperties(context, entity, dataMap, entityUrn));
8080
mappingHelper.mapToResult(
8181
INSTITUTIONAL_MEMORY_ASPECT_NAME,
8282
(mlFeatureTable, dataMap) ->
@@ -146,10 +146,10 @@ private static void mapMLFeatureTableProperties(
146146
@Nonnull DataMap dataMap,
147147
Urn entityUrn) {
148148
MLFeatureTableProperties featureTableProperties = new MLFeatureTableProperties(dataMap);
149-
mlFeatureTable.setFeatureTableProperties(
150-
MLFeatureTablePropertiesMapper.map(context, featureTableProperties, entityUrn));
151-
mlFeatureTable.setProperties(
152-
MLFeatureTablePropertiesMapper.map(context, featureTableProperties, entityUrn));
149+
com.linkedin.datahub.graphql.generated.MLFeatureTableProperties graphqlProperties =
150+
MLFeatureTablePropertiesMapper.map(context, featureTableProperties, entityUrn);
151+
mlFeatureTable.setFeatureTableProperties(graphqlProperties);
152+
mlFeatureTable.setProperties(graphqlProperties);
153153
mlFeatureTable.setDescription(featureTableProperties.getDescription());
154154
}
155155

‎datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLFeatureTablePropertiesMapper.java

+11-7
Original file line numberDiff line numberDiff line change
@@ -8,26 +8,30 @@
88
import com.linkedin.datahub.graphql.generated.MLFeatureTableProperties;
99
import com.linkedin.datahub.graphql.generated.MLPrimaryKey;
1010
import com.linkedin.datahub.graphql.types.common.mappers.CustomPropertiesMapper;
11+
import com.linkedin.datahub.graphql.types.mappers.EmbeddedModelMapper;
1112
import java.util.stream.Collectors;
13+
import javax.annotation.Nonnull;
1214
import javax.annotation.Nullable;
13-
import lombok.NonNull;
1415

15-
public class MLFeatureTablePropertiesMapper {
16+
public class MLFeatureTablePropertiesMapper
17+
implements EmbeddedModelMapper<
18+
com.linkedin.ml.metadata.MLFeatureTableProperties, MLFeatureTableProperties> {
1619

1720
public static final MLFeatureTablePropertiesMapper INSTANCE =
1821
new MLFeatureTablePropertiesMapper();
1922

2023
public static MLFeatureTableProperties map(
2124
@Nullable final QueryContext context,
22-
@NonNull final com.linkedin.ml.metadata.MLFeatureTableProperties mlFeatureTableProperties,
23-
Urn entityUrn) {
25+
@Nonnull final com.linkedin.ml.metadata.MLFeatureTableProperties mlFeatureTableProperties,
26+
@Nonnull Urn entityUrn) {
2427
return INSTANCE.apply(context, mlFeatureTableProperties, entityUrn);
2528
}
2629

27-
public static MLFeatureTableProperties apply(
30+
@Override
31+
public MLFeatureTableProperties apply(
2832
@Nullable final QueryContext context,
29-
@NonNull final com.linkedin.ml.metadata.MLFeatureTableProperties mlFeatureTableProperties,
30-
Urn entityUrn) {
33+
@Nonnull final com.linkedin.ml.metadata.MLFeatureTableProperties mlFeatureTableProperties,
34+
@Nonnull Urn entityUrn) {
3135
final MLFeatureTableProperties result = new MLFeatureTableProperties();
3236

3337
result.setDescription(mlFeatureTableProperties.getDescription());

‎datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLModelGroupMapper.java

+7-4
Original file line numberDiff line numberDiff line change
@@ -75,9 +75,8 @@ public MLModelGroup apply(
7575
mappingHelper.mapToResult(
7676
ML_MODEL_GROUP_KEY_ASPECT_NAME, MLModelGroupMapper::mapToMLModelGroupKey);
7777
mappingHelper.mapToResult(
78-
context,
7978
ML_MODEL_GROUP_PROPERTIES_ASPECT_NAME,
80-
MLModelGroupMapper::mapToMLModelGroupProperties);
79+
(entity, dataMap) -> mapToMLModelGroupProperties(context, entity, dataMap, entityUrn));
8180
mappingHelper.mapToResult(
8281
STATUS_ASPECT_NAME,
8382
(mlModelGroup, dataMap) ->
@@ -136,9 +135,13 @@ private static void mapToMLModelGroupKey(MLModelGroup mlModelGroup, DataMap data
136135
}
137136

138137
private static void mapToMLModelGroupProperties(
139-
@Nullable final QueryContext context, MLModelGroup mlModelGroup, DataMap dataMap) {
138+
@Nullable final QueryContext context,
139+
MLModelGroup mlModelGroup,
140+
DataMap dataMap,
141+
@Nonnull Urn entityUrn) {
140142
MLModelGroupProperties modelGroupProperties = new MLModelGroupProperties(dataMap);
141-
mlModelGroup.setProperties(MLModelGroupPropertiesMapper.map(context, modelGroupProperties));
143+
mlModelGroup.setProperties(
144+
MLModelGroupPropertiesMapper.map(context, modelGroupProperties, entityUrn));
142145
if (modelGroupProperties.getDescription() != null) {
143146
mlModelGroup.setDescription(modelGroupProperties.getDescription());
144147
}

‎datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLModelGroupPropertiesMapper.java

+13-6
Original file line numberDiff line numberDiff line change
@@ -1,27 +1,31 @@
11
package com.linkedin.datahub.graphql.types.mlmodel.mappers;
22

3+
import com.linkedin.common.urn.Urn;
34
import com.linkedin.datahub.graphql.QueryContext;
45
import com.linkedin.datahub.graphql.generated.MLModelGroupProperties;
5-
import com.linkedin.datahub.graphql.types.mappers.ModelMapper;
6+
import com.linkedin.datahub.graphql.types.common.mappers.CustomPropertiesMapper;
7+
import com.linkedin.datahub.graphql.types.mappers.EmbeddedModelMapper;
8+
import javax.annotation.Nonnull;
69
import javax.annotation.Nullable;
7-
import lombok.NonNull;
810

911
public class MLModelGroupPropertiesMapper
10-
implements ModelMapper<
12+
implements EmbeddedModelMapper<
1113
com.linkedin.ml.metadata.MLModelGroupProperties, MLModelGroupProperties> {
1214

1315
public static final MLModelGroupPropertiesMapper INSTANCE = new MLModelGroupPropertiesMapper();
1416

1517
public static MLModelGroupProperties map(
1618
@Nullable QueryContext context,
17-
@NonNull final com.linkedin.ml.metadata.MLModelGroupProperties mlModelGroupProperties) {
18-
return INSTANCE.apply(context, mlModelGroupProperties);
19+
@Nonnull final com.linkedin.ml.metadata.MLModelGroupProperties mlModelGroupProperties,
20+
@Nonnull Urn entityUrn) {
21+
return INSTANCE.apply(context, mlModelGroupProperties, entityUrn);
1922
}
2023

2124
@Override
2225
public MLModelGroupProperties apply(
2326
@Nullable QueryContext context,
24-
@NonNull final com.linkedin.ml.metadata.MLModelGroupProperties mlModelGroupProperties) {
27+
@Nonnull final com.linkedin.ml.metadata.MLModelGroupProperties mlModelGroupProperties,
28+
@Nonnull Urn entityUrn) {
2529
final MLModelGroupProperties result = new MLModelGroupProperties();
2630

2731
result.setDescription(mlModelGroupProperties.getDescription());
@@ -30,6 +34,9 @@ public MLModelGroupProperties apply(
3034
}
3135
result.setCreatedAt(mlModelGroupProperties.getCreatedAt());
3236

37+
result.setCustomProperties(
38+
CustomPropertiesMapper.map(mlModelGroupProperties.getCustomProperties(), entityUrn));
39+
3340
return result;
3441
}
3542
}

‎datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLModelPropertiesMapper.java

+7-5
Original file line numberDiff line numberDiff line change
@@ -7,25 +7,27 @@
77
import com.linkedin.datahub.graphql.generated.MLModelGroup;
88
import com.linkedin.datahub.graphql.generated.MLModelProperties;
99
import com.linkedin.datahub.graphql.types.common.mappers.CustomPropertiesMapper;
10+
import com.linkedin.datahub.graphql.types.mappers.EmbeddedModelMapper;
1011
import java.util.stream.Collectors;
12+
import javax.annotation.Nonnull;
1113
import javax.annotation.Nullable;
12-
import lombok.NonNull;
1314

14-
public class MLModelPropertiesMapper {
15+
public class MLModelPropertiesMapper
16+
implements EmbeddedModelMapper<com.linkedin.ml.metadata.MLModelProperties, MLModelProperties> {
1517

1618
public static final MLModelPropertiesMapper INSTANCE = new MLModelPropertiesMapper();
1719

1820
public static MLModelProperties map(
1921
@Nullable final QueryContext context,
20-
@NonNull final com.linkedin.ml.metadata.MLModelProperties mlModelProperties,
22+
@Nonnull final com.linkedin.ml.metadata.MLModelProperties mlModelProperties,
2123
Urn entityUrn) {
2224
return INSTANCE.apply(context, mlModelProperties, entityUrn);
2325
}
2426

2527
public MLModelProperties apply(
2628
@Nullable final QueryContext context,
27-
@NonNull final com.linkedin.ml.metadata.MLModelProperties mlModelProperties,
28-
Urn entityUrn) {
29+
@Nonnull final com.linkedin.ml.metadata.MLModelProperties mlModelProperties,
30+
@Nonnull Urn entityUrn) {
2931
final MLModelProperties result = new MLModelProperties();
3032

3133
result.setDate(mlModelProperties.getDate());

‎datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLPrimaryKeyMapper.java

+9-6
Original file line numberDiff line numberDiff line change
@@ -74,9 +74,8 @@ public MLPrimaryKey apply(
7474
mappingHelper.mapToResult(
7575
ML_PRIMARY_KEY_KEY_ASPECT_NAME, MLPrimaryKeyMapper::mapMLPrimaryKeyKey);
7676
mappingHelper.mapToResult(
77-
context,
7877
ML_PRIMARY_KEY_PROPERTIES_ASPECT_NAME,
79-
MLPrimaryKeyMapper::mapMLPrimaryKeyProperties);
78+
(entity, dataMap) -> mapMLPrimaryKeyProperties(context, entity, dataMap, entityUrn));
8079
mappingHelper.mapToResult(
8180
INSTITUTIONAL_MEMORY_ASPECT_NAME,
8281
(mlPrimaryKey, dataMap) ->
@@ -132,11 +131,15 @@ private static void mapMLPrimaryKeyKey(MLPrimaryKey mlPrimaryKey, DataMap dataMa
132131
}
133132

134133
private static void mapMLPrimaryKeyProperties(
135-
@Nullable final QueryContext context, MLPrimaryKey mlPrimaryKey, DataMap dataMap) {
134+
@Nullable final QueryContext context,
135+
MLPrimaryKey mlPrimaryKey,
136+
DataMap dataMap,
137+
@Nonnull Urn entityUrn) {
136138
MLPrimaryKeyProperties primaryKeyProperties = new MLPrimaryKeyProperties(dataMap);
137-
mlPrimaryKey.setPrimaryKeyProperties(
138-
MLPrimaryKeyPropertiesMapper.map(context, primaryKeyProperties));
139-
mlPrimaryKey.setProperties(MLPrimaryKeyPropertiesMapper.map(context, primaryKeyProperties));
139+
com.linkedin.datahub.graphql.generated.MLPrimaryKeyProperties graphqlProperties =
140+
MLPrimaryKeyPropertiesMapper.map(context, primaryKeyProperties, entityUrn);
141+
mlPrimaryKey.setPrimaryKeyProperties(graphqlProperties);
142+
mlPrimaryKey.setProperties(graphqlProperties);
140143
mlPrimaryKey.setDescription(primaryKeyProperties.getDescription());
141144
if (primaryKeyProperties.getDataType() != null) {
142145
mlPrimaryKey.setDataType(

‎datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLPrimaryKeyPropertiesMapper.java

+13-6
Original file line numberDiff line numberDiff line change
@@ -1,30 +1,34 @@
11
package com.linkedin.datahub.graphql.types.mlmodel.mappers;
22

3+
import com.linkedin.common.urn.Urn;
34
import com.linkedin.datahub.graphql.QueryContext;
45
import com.linkedin.datahub.graphql.generated.Dataset;
56
import com.linkedin.datahub.graphql.generated.MLFeatureDataType;
67
import com.linkedin.datahub.graphql.generated.MLPrimaryKeyProperties;
7-
import com.linkedin.datahub.graphql.types.mappers.ModelMapper;
8+
import com.linkedin.datahub.graphql.types.common.mappers.CustomPropertiesMapper;
9+
import com.linkedin.datahub.graphql.types.mappers.EmbeddedModelMapper;
810
import java.util.stream.Collectors;
11+
import javax.annotation.Nonnull;
912
import javax.annotation.Nullable;
10-
import lombok.NonNull;
1113

1214
public class MLPrimaryKeyPropertiesMapper
13-
implements ModelMapper<
15+
implements EmbeddedModelMapper<
1416
com.linkedin.ml.metadata.MLPrimaryKeyProperties, MLPrimaryKeyProperties> {
1517

1618
public static final MLPrimaryKeyPropertiesMapper INSTANCE = new MLPrimaryKeyPropertiesMapper();
1719

1820
public static MLPrimaryKeyProperties map(
1921
@Nullable QueryContext context,
20-
@NonNull final com.linkedin.ml.metadata.MLPrimaryKeyProperties mlPrimaryKeyProperties) {
21-
return INSTANCE.apply(context, mlPrimaryKeyProperties);
22+
@Nonnull final com.linkedin.ml.metadata.MLPrimaryKeyProperties mlPrimaryKeyProperties,
23+
@Nonnull Urn entityUrn) {
24+
return INSTANCE.apply(context, mlPrimaryKeyProperties, entityUrn);
2225
}
2326

2427
@Override
2528
public MLPrimaryKeyProperties apply(
2629
@Nullable QueryContext context,
27-
@NonNull final com.linkedin.ml.metadata.MLPrimaryKeyProperties mlPrimaryKeyProperties) {
30+
@Nonnull final com.linkedin.ml.metadata.MLPrimaryKeyProperties mlPrimaryKeyProperties,
31+
@Nonnull Urn entityUrn) {
2832
final MLPrimaryKeyProperties result = new MLPrimaryKeyProperties();
2933

3034
result.setDescription(mlPrimaryKeyProperties.getDescription());
@@ -45,6 +49,9 @@ public MLPrimaryKeyProperties apply(
4549
})
4650
.collect(Collectors.toList()));
4751

52+
result.setCustomProperties(
53+
CustomPropertiesMapper.map(mlPrimaryKeyProperties.getCustomProperties(), entityUrn));
54+
4855
return result;
4956
}
5057
}

‎datahub-graphql-core/src/main/resources/entity.graphql

+9-3
Original file line numberDiff line numberDiff line change
@@ -9829,11 +9829,13 @@ type MLModelGroup implements EntityWithRelationships & Entity & BrowsableEntity
98299829

98309830
type MLModelGroupProperties {
98319831

9832-
description: String
9832+
description: String
98339833

98349834
createdAt: Long
98359835

98369836
version: VersionTag
9837+
9838+
customProperties: [CustomPropertiesEntry!]
98379839
}
98389840

98399841
"""
@@ -10028,6 +10030,8 @@ type MLFeatureProperties {
1002810030
version: VersionTag
1002910031

1003010032
sources: [Dataset]
10033+
10034+
customProperties: [CustomPropertiesEntry!]
1003110035
}
1003210036

1003310037
"""
@@ -10164,13 +10168,15 @@ type MLPrimaryKey implements EntityWithRelationships & Entity {
1016410168

1016510169
type MLPrimaryKeyProperties {
1016610170

10167-
description: String
10171+
description: String
1016810172

1016910173
dataType: MLFeatureDataType
1017010174

1017110175
version: VersionTag
1017210176

1017310177
sources: [Dataset]
10178+
10179+
customProperties: [CustomPropertiesEntry!]
1017410180
}
1017510181

1017610182
"""
@@ -10347,7 +10353,7 @@ type MLModelGroupEditableProperties {
1034710353

1034810354
type MLFeatureTableProperties {
1034910355

10350-
description: String
10356+
description: String
1035110357

1035210358
mlFeatures: [MLFeature]
1035310359

0 commit comments

Comments
 (0)