Problem

When training a model using data stored in a table in Unity Catalog, the lineage to the upstream dataset(s) is tracked using mlflow.log_input, which logs the input table with the MLflow run. However, lineage for the output table (containing predictions) is not tracked.

Cause

There is no built-in method in MLflow to log output tables, similar to mlflow.log_input for input tables. 

Solution

Save the output table as a CSV file and log it as an artifact. This way, you can indirectly track the lineage of the output table. You can use the following code. 

import mlflow from sklearn import datasets from sklearn.ensemble import RandomForestRegressor import pandas as pd import tempfile import os # Load dataset dataset = mlflow.data.load_delta(table_name="<your-catalog>.<your-schema>.<your-table-name>", version="0") pd_df = dataset.df.toPandas() X = pd_df.drop("species", axis=1) y = pd_df["species"] # Train model and log input table with mlflow.start_run() as run:     clf = RandomForestRegressor(n_estimators=100)     clf.fit(X, y)     mlflow.log_input(dataset, "training")     # Make predictions     predictions = clf.predict(X)     pd_df["predictions"] = predictions     # Save predictions to an output table     <your-output-table-name> = "<your-catalog>.<your-schema>.<your-iris-output>"     output_df = spark.createDataFrame(pd_df) output_df.write.format("delta").mode("overwrite").saveAsTable(<your-output-table-name>)     # Log the output table     with tempfile.TemporaryDirectory() as tmpdir:         temp_path = os.path.join(tmpdir, "predictions.csv")         pd_df.to_csv(temp_path, index=False)         # Log the temporary file as an artifact         mlflow.log_artifact(temp_path, "output_table") print(f"Output table {<your-output-table-name>} created successfully.")

To locate the output table, navigate to the Artifacts page of the specific MLflow run. 

By following the above steps, we can ensure that the output table can be traced back to the corresponding run.