yujunjun/datahub
2
0
Fork 0
mirror of https://github.com/datahub-project/datahub.git synced 2025-12-12 10:35:51 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(graphql/ml): Add custom properties to ml entities (#12152)

Browse Source
This commit is contained in:
Andrew Sikowitz 2024-12-19 10:20:06 -08:00 committed by GitHub
parent 9031b49b23
commit b7bb5ca7ee
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
11 changed files with 108 additions and 52 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mappers/EmbeddedModelMapper.java Normal file
View File

@ -0,0 +1,12 @@
package com.linkedin.datahub.graphql.types.mappers;
import com.linkedin.common.urn.Urn;
import com.linkedin.datahub.graphql.QueryContext;
import javax.annotation.Nonnull;
import javax.annotation.Nullable;
/** Made for models that are embedded in other models and thus do not encode their own URN. */
public interface EmbeddedModelMapper<I, O> {
O apply(
@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
View File

@ -75,7 +75,8 @@ public class MLFeatureMapper implements ModelMapper<EntityResponse, MLFeature> {
mlFeature.setOwnership(
OwnershipMapper.map(context, new Ownership(dataMap), entityUrn)));
mappingHelper.mapToResult(
context, ML_FEATURE_PROPERTIES_ASPECT_NAME, MLFeatureMapper::mapMLFeatureProperties);
ML_FEATURE_PROPERTIES_ASPECT_NAME,
(entity, dataMap) -> mapMLFeatureProperties(context, entity, dataMap, entityUrn));
mappingHelper.mapToResult(
INSTITUTIONAL_MEMORY_ASPECT_NAME,
(mlFeature, dataMap) ->
@ -138,10 +139,13 @@ public class MLFeatureMapper implements ModelMapper<EntityResponse, MLFeature> {
private static void mapMLFeatureProperties(
@Nullable final QueryContext context,
@Nonnull MLFeature mlFeature,
@Nonnull DataMap dataMap) {
@Nonnull DataMap dataMap,
@Nonnull Urn entityUrn) {
MLFeatureProperties featureProperties = new MLFeatureProperties(dataMap);
mlFeature.setFeatureProperties(MLFeaturePropertiesMapper.map(context, featureProperties));
mlFeature.setProperties(MLFeaturePropertiesMapper.map(context, featureProperties));
com.linkedin.datahub.graphql.generated.MLFeatureProperties graphqlProperties =
MLFeaturePropertiesMapper.map(context, featureProperties, entityUrn);
mlFeature.setFeatureProperties(graphqlProperties);
mlFeature.setProperties(graphqlProperties);
mlFeature.setDescription(featureProperties.getDescription());
if (featureProperties.getDataType() != null) {
mlFeature.setDataType(MLFeatureDataType.valueOf(featureProperties.getDataType().toString()));

20
datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLFeaturePropertiesMapper.java
View File

@ -1,29 +1,34 @@
package com.linkedin.datahub.graphql.types.mlmodel.mappers;
import com.linkedin.common.urn.Urn;
import com.linkedin.datahub.graphql.QueryContext;
import com.linkedin.datahub.graphql.generated.Dataset;
import com.linkedin.datahub.graphql.generated.MLFeatureDataType;
import com.linkedin.datahub.graphql.generated.MLFeatureProperties;
import com.linkedin.datahub.graphql.types.mappers.ModelMapper;
import com.linkedin.datahub.graphql.types.common.mappers.CustomPropertiesMapper;
import com.linkedin.datahub.graphql.types.mappers.EmbeddedModelMapper;
import java.util.stream.Collectors;
import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import lombok.NonNull;
public class MLFeaturePropertiesMapper
implements ModelMapper<com.linkedin.ml.metadata.MLFeatureProperties, MLFeatureProperties> {
implements EmbeddedModelMapper<
com.linkedin.ml.metadata.MLFeatureProperties, MLFeatureProperties> {
public static final MLFeaturePropertiesMapper INSTANCE = new MLFeaturePropertiesMapper();
public static MLFeatureProperties map(
@Nullable QueryContext context,
@NonNull final com.linkedin.ml.metadata.MLFeatureProperties mlFeatureProperties) {
return INSTANCE.apply(context, mlFeatureProperties);
@Nonnull final com.linkedin.ml.metadata.MLFeatureProperties mlFeatureProperties,
@Nonnull Urn entityUrn) {
return INSTANCE.apply(context, mlFeatureProperties, entityUrn);
}
@Override
public MLFeatureProperties apply(
@Nullable QueryContext context,
@NonNull final com.linkedin.ml.metadata.MLFeatureProperties mlFeatureProperties) {
@Nonnull final com.linkedin.ml.metadata.MLFeatureProperties mlFeatureProperties,
@Nonnull Urn entityUrn) {
final MLFeatureProperties result = new MLFeatureProperties();
result.setDescription(mlFeatureProperties.getDescription());
@ -45,6 +50,9 @@ public class MLFeaturePropertiesMapper
.collect(Collectors.toList()));
}
result.setCustomProperties(
CustomPropertiesMapper.map(mlFeatureProperties.getCustomProperties(), entityUrn));
return result;
}
}

10
datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLFeatureTableMapper.java
View File

@ -76,7 +76,7 @@ public class MLFeatureTableMapper implements ModelMapper<EntityResponse, MLFeatu
mappingHelper.mapToResult(ML_FEATURE_TABLE_KEY_ASPECT_NAME, this::mapMLFeatureTableKey);
mappingHelper.mapToResult(
ML_FEATURE_TABLE_PROPERTIES_ASPECT_NAME,
(entity, dataMap) -> this.mapMLFeatureTableProperties(context, entity, dataMap, entityUrn));
(entity, dataMap) -> mapMLFeatureTableProperties(context, entity, dataMap, entityUrn));
mappingHelper.mapToResult(
INSTITUTIONAL_MEMORY_ASPECT_NAME,
(mlFeatureTable, dataMap) ->
@ -146,10 +146,10 @@ public class MLFeatureTableMapper implements ModelMapper<EntityResponse, MLFeatu
@Nonnull DataMap dataMap,
Urn entityUrn) {
MLFeatureTableProperties featureTableProperties = new MLFeatureTableProperties(dataMap);
mlFeatureTable.setFeatureTableProperties(
MLFeatureTablePropertiesMapper.map(context, featureTableProperties, entityUrn));
mlFeatureTable.setProperties(
MLFeatureTablePropertiesMapper.map(context, featureTableProperties, entityUrn));
com.linkedin.datahub.graphql.generated.MLFeatureTableProperties graphqlProperties =
MLFeatureTablePropertiesMapper.map(context, featureTableProperties, entityUrn);
mlFeatureTable.setFeatureTableProperties(graphqlProperties);
mlFeatureTable.setProperties(graphqlProperties);
mlFeatureTable.setDescription(featureTableProperties.getDescription());
}

18
datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLFeatureTablePropertiesMapper.java
View File

@ -8,26 +8,30 @@ import com.linkedin.datahub.graphql.generated.MLFeature;
import com.linkedin.datahub.graphql.generated.MLFeatureTableProperties;
import com.linkedin.datahub.graphql.generated.MLPrimaryKey;
import com.linkedin.datahub.graphql.types.common.mappers.CustomPropertiesMapper;
import com.linkedin.datahub.graphql.types.mappers.EmbeddedModelMapper;
import java.util.stream.Collectors;
import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import lombok.NonNull;
public class MLFeatureTablePropertiesMapper {
public class MLFeatureTablePropertiesMapper
implements EmbeddedModelMapper<
com.linkedin.ml.metadata.MLFeatureTableProperties, MLFeatureTableProperties> {
public static final MLFeatureTablePropertiesMapper INSTANCE =
new MLFeatureTablePropertiesMapper();
public static MLFeatureTableProperties map(
@Nullable final QueryContext context,
@NonNull final com.linkedin.ml.metadata.MLFeatureTableProperties mlFeatureTableProperties,
Urn entityUrn) {
@Nonnull final com.linkedin.ml.metadata.MLFeatureTableProperties mlFeatureTableProperties,
@Nonnull Urn entityUrn) {
return INSTANCE.apply(context, mlFeatureTableProperties, entityUrn);
}
public static MLFeatureTableProperties apply(
@Override
public MLFeatureTableProperties apply(
@Nullable final QueryContext context,
@NonNull final com.linkedin.ml.metadata.MLFeatureTableProperties mlFeatureTableProperties,
Urn entityUrn) {
@Nonnull final com.linkedin.ml.metadata.MLFeatureTableProperties mlFeatureTableProperties,
@Nonnull Urn entityUrn) {
final MLFeatureTableProperties result = new MLFeatureTableProperties();
result.setDescription(mlFeatureTableProperties.getDescription());

11
datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLModelGroupMapper.java
View File

@ -75,9 +75,8 @@ public class MLModelGroupMapper implements ModelMapper<EntityResponse, MLModelGr
mappingHelper.mapToResult(
ML_MODEL_GROUP_KEY_ASPECT_NAME, MLModelGroupMapper::mapToMLModelGroupKey);
mappingHelper.mapToResult(
context,
ML_MODEL_GROUP_PROPERTIES_ASPECT_NAME,
MLModelGroupMapper::mapToMLModelGroupProperties);
(entity, dataMap) -> mapToMLModelGroupProperties(context, entity, dataMap, entityUrn));
mappingHelper.mapToResult(
STATUS_ASPECT_NAME,
(mlModelGroup, dataMap) ->
@ -136,9 +135,13 @@ public class MLModelGroupMapper implements ModelMapper<EntityResponse, MLModelGr
}
private static void mapToMLModelGroupProperties(
@Nullable final QueryContext context, MLModelGroup mlModelGroup, DataMap dataMap) {
@Nullable final QueryContext context,
MLModelGroup mlModelGroup,
DataMap dataMap,
@Nonnull Urn entityUrn) {
MLModelGroupProperties modelGroupProperties = new MLModelGroupProperties(dataMap);
mlModelGroup.setProperties(MLModelGroupPropertiesMapper.map(context, modelGroupProperties));
mlModelGroup.setProperties(
MLModelGroupPropertiesMapper.map(context, modelGroupProperties, entityUrn));
if (modelGroupProperties.getDescription() != null) {
mlModelGroup.setDescription(modelGroupProperties.getDescription());
}

19
datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLModelGroupPropertiesMapper.java
View File

@ -1,27 +1,31 @@
package com.linkedin.datahub.graphql.types.mlmodel.mappers;
import com.linkedin.common.urn.Urn;
import com.linkedin.datahub.graphql.QueryContext;
import com.linkedin.datahub.graphql.generated.MLModelGroupProperties;
import com.linkedin.datahub.graphql.types.mappers.ModelMapper;
import com.linkedin.datahub.graphql.types.common.mappers.CustomPropertiesMapper;
import com.linkedin.datahub.graphql.types.mappers.EmbeddedModelMapper;
import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import lombok.NonNull;
public class MLModelGroupPropertiesMapper
implements ModelMapper<
implements EmbeddedModelMapper<
com.linkedin.ml.metadata.MLModelGroupProperties, MLModelGroupProperties> {
public static final MLModelGroupPropertiesMapper INSTANCE = new MLModelGroupPropertiesMapper();
public static MLModelGroupProperties map(
@Nullable QueryContext context,
@NonNull final com.linkedin.ml.metadata.MLModelGroupProperties mlModelGroupProperties) {
return INSTANCE.apply(context, mlModelGroupProperties);
@Nonnull final com.linkedin.ml.metadata.MLModelGroupProperties mlModelGroupProperties,
@Nonnull Urn entityUrn) {
return INSTANCE.apply(context, mlModelGroupProperties, entityUrn);
}
@Override
public MLModelGroupProperties apply(
@Nullable QueryContext context,
@NonNull final com.linkedin.ml.metadata.MLModelGroupProperties mlModelGroupProperties) {
@Nonnull final com.linkedin.ml.metadata.MLModelGroupProperties mlModelGroupProperties,
@Nonnull Urn entityUrn) {
final MLModelGroupProperties result = new MLModelGroupProperties();
result.setDescription(mlModelGroupProperties.getDescription());
@ -30,6 +34,9 @@ public class MLModelGroupPropertiesMapper
}
result.setCreatedAt(mlModelGroupProperties.getCreatedAt());
result.setCustomProperties(
CustomPropertiesMapper.map(mlModelGroupProperties.getCustomProperties(), entityUrn));
return result;
}
}

12
datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLModelPropertiesMapper.java
View File

@ -7,25 +7,27 @@ import com.linkedin.datahub.graphql.QueryContext;
import com.linkedin.datahub.graphql.generated.MLModelGroup;
import com.linkedin.datahub.graphql.generated.MLModelProperties;
import com.linkedin.datahub.graphql.types.common.mappers.CustomPropertiesMapper;
import com.linkedin.datahub.graphql.types.mappers.EmbeddedModelMapper;
import java.util.stream.Collectors;
import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import lombok.NonNull;
public class MLModelPropertiesMapper {
public class MLModelPropertiesMapper
implements EmbeddedModelMapper<com.linkedin.ml.metadata.MLModelProperties, MLModelProperties> {
public static final MLModelPropertiesMapper INSTANCE = new MLModelPropertiesMapper();
public static MLModelProperties map(
@Nullable final QueryContext context,
@NonNull final com.linkedin.ml.metadata.MLModelProperties mlModelProperties,
@Nonnull final com.linkedin.ml.metadata.MLModelProperties mlModelProperties,
Urn entityUrn) {
return INSTANCE.apply(context, mlModelProperties, entityUrn);
}
public MLModelProperties apply(
@Nullable final QueryContext context,
@NonNull final com.linkedin.ml.metadata.MLModelProperties mlModelProperties,
Urn entityUrn) {
@Nonnull final com.linkedin.ml.metadata.MLModelProperties mlModelProperties,
@Nonnull Urn entityUrn) {
final MLModelProperties result = new MLModelProperties();
result.setDate(mlModelProperties.getDate());

15
datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLPrimaryKeyMapper.java
View File

@ -74,9 +74,8 @@ public class MLPrimaryKeyMapper implements ModelMapper<EntityResponse, MLPrimary
mappingHelper.mapToResult(
ML_PRIMARY_KEY_KEY_ASPECT_NAME, MLPrimaryKeyMapper::mapMLPrimaryKeyKey);
mappingHelper.mapToResult(
context,
ML_PRIMARY_KEY_PROPERTIES_ASPECT_NAME,
MLPrimaryKeyMapper::mapMLPrimaryKeyProperties);
(entity, dataMap) -> mapMLPrimaryKeyProperties(context, entity, dataMap, entityUrn));
mappingHelper.mapToResult(
INSTITUTIONAL_MEMORY_ASPECT_NAME,
(mlPrimaryKey, dataMap) ->
@ -132,11 +131,15 @@ public class MLPrimaryKeyMapper implements ModelMapper<EntityResponse, MLPrimary
}
private static void mapMLPrimaryKeyProperties(
@Nullable final QueryContext context, MLPrimaryKey mlPrimaryKey, DataMap dataMap) {
@Nullable final QueryContext context,
MLPrimaryKey mlPrimaryKey,
DataMap dataMap,
@Nonnull Urn entityUrn) {
MLPrimaryKeyProperties primaryKeyProperties = new MLPrimaryKeyProperties(dataMap);
mlPrimaryKey.setPrimaryKeyProperties(
MLPrimaryKeyPropertiesMapper.map(context, primaryKeyProperties));
mlPrimaryKey.setProperties(MLPrimaryKeyPropertiesMapper.map(context, primaryKeyProperties));
com.linkedin.datahub.graphql.generated.MLPrimaryKeyProperties graphqlProperties =
MLPrimaryKeyPropertiesMapper.map(context, primaryKeyProperties, entityUrn);
mlPrimaryKey.setPrimaryKeyProperties(graphqlProperties);
mlPrimaryKey.setProperties(graphqlProperties);
mlPrimaryKey.setDescription(primaryKeyProperties.getDescription());
if (primaryKeyProperties.getDataType() != null) {
mlPrimaryKey.setDataType(

19
datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLPrimaryKeyPropertiesMapper.java
View File

@ -1,30 +1,34 @@
package com.linkedin.datahub.graphql.types.mlmodel.mappers;
import com.linkedin.common.urn.Urn;
import com.linkedin.datahub.graphql.QueryContext;
import com.linkedin.datahub.graphql.generated.Dataset;
import com.linkedin.datahub.graphql.generated.MLFeatureDataType;
import com.linkedin.datahub.graphql.generated.MLPrimaryKeyProperties;
import com.linkedin.datahub.graphql.types.mappers.ModelMapper;
import com.linkedin.datahub.graphql.types.common.mappers.CustomPropertiesMapper;
import com.linkedin.datahub.graphql.types.mappers.EmbeddedModelMapper;
import java.util.stream.Collectors;
import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import lombok.NonNull;
public class MLPrimaryKeyPropertiesMapper
implements ModelMapper<
implements EmbeddedModelMapper<
com.linkedin.ml.metadata.MLPrimaryKeyProperties, MLPrimaryKeyProperties> {
public static final MLPrimaryKeyPropertiesMapper INSTANCE = new MLPrimaryKeyPropertiesMapper();
public static MLPrimaryKeyProperties map(
@Nullable QueryContext context,
@NonNull final com.linkedin.ml.metadata.MLPrimaryKeyProperties mlPrimaryKeyProperties) {
return INSTANCE.apply(context, mlPrimaryKeyProperties);
@Nonnull final com.linkedin.ml.metadata.MLPrimaryKeyProperties mlPrimaryKeyProperties,
@Nonnull Urn entityUrn) {
return INSTANCE.apply(context, mlPrimaryKeyProperties, entityUrn);
}
@Override
public MLPrimaryKeyProperties apply(
@Nullable QueryContext context,
@NonNull final com.linkedin.ml.metadata.MLPrimaryKeyProperties mlPrimaryKeyProperties) {
@Nonnull final com.linkedin.ml.metadata.MLPrimaryKeyProperties mlPrimaryKeyProperties,
@Nonnull Urn entityUrn) {
final MLPrimaryKeyProperties result = new MLPrimaryKeyProperties();
result.setDescription(mlPrimaryKeyProperties.getDescription());
@ -45,6 +49,9 @@ public class MLPrimaryKeyPropertiesMapper
})
.collect(Collectors.toList()));
result.setCustomProperties(
CustomPropertiesMapper.map(mlPrimaryKeyProperties.getCustomProperties(), entityUrn));
return result;
}
}

12
datahub-graphql-core/src/main/resources/entity.graphql
View File

@ -9829,11 +9829,13 @@ type MLModelGroup implements EntityWithRelationships & Entity & BrowsableEntity
type MLModelGroupProperties {
description: String
description: String
createdAt: Long
version: VersionTag
customProperties: [CustomPropertiesEntry!]
}
"""
@ -10028,6 +10030,8 @@ type MLFeatureProperties {
version: VersionTag
sources: [Dataset]
customProperties: [CustomPropertiesEntry!]
}
"""
@ -10164,13 +10168,15 @@ type MLPrimaryKey implements EntityWithRelationships & Entity {
type MLPrimaryKeyProperties {
description: String
description: String
dataType: MLFeatureDataType
version: VersionTag
sources: [Dataset]
customProperties: [CustomPropertiesEntry!]
}
"""
@ -10347,7 +10353,7 @@ type MLModelGroupEditableProperties {
type MLFeatureTableProperties {
description: String
description: String
mlFeatures: [MLFeature]