import mlflow

mlflow.set_registry_uri("databricks-uc")
model_uri=f"models:/{model_name}@DLT" 
loaded_model_udf = mlflow.pyfunc.spark_udf(spark, model_uri=model_uri, result_type="string")

primary_key = "customerID"
features = ["SeniorCitizen", "tenure", "MonthlyCharges", "TotalCharges"]

import dlt
from pyspark.sql.functions import col, struct

@dlt.table(
  name="raw_inputs",
  comment="Raw inputs table",
  table_properties={
    "quality": "bronze"
  } 
)
def raw_inputs():
  return spark.read.csv(bronze_dataset_path, inferSchema=True, header=True, multiLine=True, escape='"')
  
@dlt.table(
  name="features_input",
  comment="Features table",
  table_properties={
    "quality": "silver"
  }
)
def features_input():
  return (
    dlt.read("raw_inputs")
    .select(primary_key, *features)
    .withColumn("SeniorCitizen",col("SeniorCitizen").cast('double'))
    .withColumn("tenure",col("tenure").cast('double'))
    .withColumn("TotalCharges",col("TotalCharges").cast('double'))
    .na.drop(how='any')
  )

@dlt.table(
  name="model_predictions",
  comment="Inference table",
  table_properties={
    "quality": "gold"
  }
)
def model_predictions():
  return (
    dlt.read("features_input")
    .withColumn("prediction", loaded_model_udf(struct(features)))
  )