Skip to content

Add ml-lineage datapack: warehouse to production ML models with column-level lineage - #224

Open
malharlakdawala wants to merge 1 commit into
datahub-project:mainfrom
malharlakdawala:add-ml-lineage-datapack
Open

malharlakdawala wants to merge 1 commit into
datahub-project:mainfrom
malharlakdawala:add-ml-lineage-datapack

Conversation

@malharlakdawala

Copy link
Copy Markdown

Summary

Adds an ml-lineage datapack: a small NYC-taxi-shaped warehouse wired all the way through
to production ML models, so the ML half of the metadata model has sample data to exercise.

Neither shipped datapack currently contains ML entities. bootstrap is datasets, dashboards,
users and tags; showcase-ecommerce is a rich BI/governance graph across Snowflake, Looker,
PowerBI and Tableau. Nothing seeds mlModel, mlModelGroup, mlFeatureTable or mlFeature,
which makes the ML lineage features and the DerivedFrom / Consumes / MemberOf
relationships awkward to demo or test against.

What it contains

25 entities, 32 aspects, 68 KB.

7 Snowflake datasets raw → staging → feature tables, with schemaMetadata
21 column-level mappings upstreamLineage.fineGrainedLineages, FIELD_SET → FIELD
2 mlFeatureTable + 7 mlFeature Feast-style, sources pointing at the feature datasets
3 mlModel in 3 mlModelGroup MLflow-style, mlFeatures wired to the features

The graph is deliberately shaped so column-level lineage is discriminating rather than
uniform — some raw columns reach a model and some do not:

raw_payments.tip_amount    -> trip_features.tip_rate, tip_variance  -> tip_predictor_v3
raw_payments.payment_type  -> stg_payments.payment_type             -> (no model)
raw_trips.trip_distance    -> distance_km, fare_per_mile            -> all three models

That makes it useful for testing impact analysis: a table-level query from raw_payments
reaches all three models, while the column-level edges reach one, none, or all depending on
which column you start from.

One feature, tip_volatility, is intentionally not named after its source column
(trip_features.tip_variance), because feature stores rename and a fixture where every
feature name matches its column hides that case.

Format

Follows the showcase-ecommerce layout: index.json with a version and an ordered files
array, flat file paths, wait_for_completion on the two stages later files depend on.
Loads with datahub datapack load ml-lineage once registered.

Registry entry

Registration lives in the main repo, so this PR only adds the content. Happy to open the
matching registry.json PR against datahub-project/datahub once this merges and the raw
URL resolves — flagging rather than opening it now so the registry never points at a 404:

"ml-lineage": {
  "name": "ml-lineage",
  "description": "ML lineage demo: NYC-taxi warehouse with column-level lineage through feature tables to production models",
  "url": "https://raw.githubusercontent.com/datahub-project/static-assets/main/datapacks/ml-lineage/index.json",
  "size_hint": "~68 KB",
  "tags": ["demo", "ml", "mlflow", "feast", "lineage", "column-level-lineage"],
  "trust": "verified",
  "pack_format_version": "1"
}

Testing

Loaded into a local DataHub Core v1.7.0 quickstart. All 25 entities resolve in the UI, the
column-level lineage renders in the dataset lineage column view, and GraphQL
searchAcrossLineage traverses dataset → mlFeature → mlModel → mlModelGroup with the
expected degree values (1 through 4) and full paths.

25 entities: 7 Snowflake datasets with 21 column-level lineage mappings,
2 mlFeatureTables with 7 mlFeatures, and 3 mlModels in 3 mlModelGroups.

Neither bootstrap nor showcase-ecommerce contains ML entities, so there is
no sample data for mlModel/mlFeature lineage or the DerivedFrom/Consumes/
MemberOf relationships. The graph is shaped so column-level lineage is
discriminating: different raw columns reach one, none, or all three models.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant