99_init.py

# Databricks notebook source
dbutils.widgets.text("reset_all_data", "false", "Reset Data")
reset_all_data = dbutils.widgets.get("reset_all_data") == "true"

# COMMAND ----------

from pyspark.sql.functions import pandas_udf
import pandas as pd
import pyspark.sql.functions as F
from pyspark.sql.functions import col, udf, length, pandas_udf
import os
import mlflow
import yaml
from typing import Iterator
from mlflow import MlflowClient
mlflow.set_registry_uri('databricks-uc')

# Set up logging
import logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logging.getLogger("py4j.java_gateway").setLevel(logging.ERROR)
logging.getLogger("py4j.clientserver").setLevel(logging.ERROR)
logging.getLogger('mlflow').setLevel(logging.ERROR) # Disable MLflow warnings
from urllib3.connectionpool import log as urllib3_log
urllib3_log.setLevel(logging.ERROR)

# Workaround for a bug fix that is in progress
mlflow.spark.autolog(disable=True)

import warnings
warnings.filterwarnings("ignore")

# COMMAND ----------

if reset_all_data:
  print(f'clearing up schema {config.CATALOG}.{config.SCHEMA}')
  _ = spark.sql(f"DROP DATABASE IF EXISTS `{config.CATALOG}.{config.SCHEMA}` CASCADE")

# COMMAND ----------

def use_and_create_db(CATALOG, SCHEMA, cloud_storage_path = None):
  print(f"USE CATALOG `{CATALOG}`")
  _ = spark.sql(f"USE CATALOG `{CATALOG}`")
  _ = spark.sql(f"""CREATE DATABASE IF NOT EXISTS `{SCHEMA}` """)

#If the catalog is defined, we force it to the given value and throw exception if not.
if len(config.CATALOG) > 0:
  current_catalog = spark.sql("SELECT current_catalog()").collect()[0]['current_catalog()']
  if current_catalog != config.CATALOG:
    catalogs = [r['catalog'] for r in spark.sql("SHOW CATALOGS").collect()]
    if config.CATALOG not in catalogs:
      _ = spark.sql(f"CREATE CATALOG IF NOT EXISTS {config.CATALOG}")
  use_and_create_db(config.CATALOG, config.SCHEMA)

print(f"using catalog.database `{config.CATALOG}`.`{config.SCHEMA}`")
_ = spark.sql(f"""USE `{config.CATALOG}`.`{config.SCHEMA}`""")    

# COMMAND ----------

if not spark.catalog.tableExists(config.SOURCE_TABLE_FULLNAME) or spark.table(config.SOURCE_TABLE_FULLNAME).isEmpty() or \
    not spark.catalog.tableExists(config.EVALUATION_TABLE_FULLNAME) or spark.table(config.EVALUATION_TABLE_FULLNAME).isEmpty():
  _ = spark.sql(f'''CREATE TABLE IF NOT EXISTS {config.SOURCE_TABLE_FULLNAME} (
            id BIGINT GENERATED BY DEFAULT AS IDENTITY,
            url STRING,
            content STRING
          ) TBLPROPERTIES (delta.enableChangeDataFeed = true)''')
  (spark.createDataFrame(pd.read_parquet('https://notebooks.databricks.com/demos/dbdemos-dataset/llm/databricks-documentation/databricks_documentation.parquet'))
   .drop('title').write.mode('overwrite').saveAsTable(config.SOURCE_TABLE_FULLNAME))
  (spark.createDataFrame(pd.read_parquet('https://notebooks.databricks.com/demos/dbdemos-dataset/llm/databricks-documentation/databricks_doc_eval_set.parquet'))
   .write.mode('overwrite').saveAsTable(config.EVALUATION_TABLE_FULLNAME))
  # Make sure enableChangeDataFeed is enabled
  _ = spark.sql('ALTER TABLE databricks_documentation SET TBLPROPERTIES (delta.enableChangeDataFeed = true)')