cdm_data_loaders.utils.spark ¶

def get_spark(
    app_name: str | None = None,
    **kwargs,
) -> SparkSession:
    """Set up a Spark session and namespace.

    If no app_name is provided, the function will use the string "cdm_data_loader".

    :param app_name: name for the app (used for the Spark session), defaults to None
    :type app_name: str, optional
    :param kwargs: extra arguments to supply to the spark session
    :type kwargs: dict[str, Any] | None
    :return: a Spark session
    :rtype: SparkSession
    """
    from berdl_notebook_utils.setup_spark_session import get_spark_session  # noqa: PLC0415

    return get_spark_session(app_name or DEFAULT_APP_NAME, **kwargs)

def set_up_workspace(
    app_name: str | None = None,
    namespace: str | None = None,
    tenant_name: str | None = None,
    **kwargs: str | int | bool | dict | None,
) -> tuple[SparkSession, str]:
    """Set up a Spark session and namespace.

    If no app_name is provided, the function will use the string "cdm_data_loader".

    N.b. this is the only function where the "raw" namespace and tenant are specified. All other functions use
    the catalog_db namespace generated by notebook_utils (from parent spark_notebook docker container)

    :param app_name: name for the app (used for the Spark session), defaults to None
    :type app_name: str, optional
    :param namespace: namespace for data I/O, defaults to None
    :type namespace: str, optional
    :param tenant_name: name of the tenant, if applicable; defaults to None
    :type tenant_name: str, optional
    :param kwargs: extra arguments to supply to the spark session
    :type kwargs: dict[str, Any] | None
    :return: a Spark session and the BERDL-approved namespace
    :rtype: tuple[SparkSession, str]
    """
    from berdl_notebook_utils.spark.database import create_namespace_if_not_exists  # noqa: PLC0415

    spark = get_spark(app_name or DEFAULT_APP_NAME, **kwargs)
    # if the namespace has already been created, this will return the doctored namespace
    catalog_db = create_namespace_if_not_exists(spark, namespace or DEFAULT_NAMESPACE, tenant_name=tenant_name)
    return (spark, catalog_db)

def get_existing_database_save_dir(spark: SparkSession, catalog_db: str) -> str | None:
    """Check whether a table is already registered in the Hive metastore, and if so, where data is saved.

    :param spark: spark session
    :type spark: SparkSession
    :param catalog_db: catalog_db name
    :type catalog_db: str
    :raises ValueError: if the table exists but the SQL query to retrieve the description returns unexpected results
    :return: path to the existing table data or None
    :rtype: str | None
    """
    if not spark.catalog.databaseExists(catalog_db):
        return None

    # namespace query
    ns_location_rows = spark.sql(f"DESCRIBE NAMESPACE EXTENDED {catalog_db}").collect()
    # we want info_value
    ns_locations = [r["info_value"] for r in ns_location_rows if r["info_name"].lower() == "location"]

    if len(ns_locations) == 1:
        return ns_locations[0]

    # there should only be one 'location' entry, so more than one indicates some sort of dodgy data.
    msg = f"Expected 1 location row in {catalog_db} description; got {len(ns_locations)} rows."
    logger.error(msg)
    logger.error(ns_locations)
    raise ValueError(msg)

def get_existing_table_save_dir(spark: SparkSession, catalog_db: str, table: str | None = None) -> str | None:
    """Check whether a table is already registered in the Hive metastore, and if so, where data is saved.

    :param spark: spark session
    :type spark: SparkSession
    :param catalog_db_table: catalog_db.table name
    :type catalog_db_table: str
    :raises ValueError: if the table exists but the SQL query to retrieve the description returns unexpected results
    :return: path to the existing table data or None
    :rtype: str | None
    """
    catalog_db_table = f"{catalog_db}.{table}"
    if not spark.catalog.tableExists(catalog_db_table):
        return None

    table_location_rows = spark.sql(f"DESCRIBE EXTENDED {catalog_db_table}").where('col_name = "Location"').collect()
    table_locations = [r["data_type"] for r in table_location_rows if r["col_name"].lower() == "location"]
    if len(table_locations) == 1:
        return table_locations[0]

    # there should only be one 'location' entry, so more than one indicates some sort of dodgy data.
    msg = f"Expected 1 location row in {catalog_db_table} description; got {len(table_locations)} rows."
    logger.error(msg)
    logger.error(table_locations)
    raise ValueError(msg)

def write_table(
    spark: SparkSession, sdf: DataFrame, catalog_db: str, table: str, mode: str = DEFAULT_WRITE_MODE
) -> None:
    """Write data as database tables and register it in the catalog.

    :param spark: spark sesh
    :type spark: SparkSession
    :param sdf: data frame to be written
    :type sdf: DataFrame
    :param catalog_db: catalog database name assigned by `set_up_workspace` or `create_namespace_if_not_exists`
    :type catalog_db: str
    :param table: table name
    :type table: str
    :param mode: spark write mode; one of the modes specified in WRITE_MODE
    :type mode: str
    """
    catalog_db_table = f"{catalog_db}.{table}"
    if mode not in WRITE_MODE:
        msg = f"Invalid mode supplied for writing table {catalog_db_table}: {mode}"
        logger.error(msg)
        raise ValueError(msg)

    if sdf is None or not isinstance(sdf, DataFrame) or sdf.isEmpty():
        logger.warning("No data to write to %s", catalog_db_table)
        return

    # this should have been set using create_namespace_if_not_exists
    base_path = get_existing_database_save_dir(spark, catalog_db)
    if not base_path:
        msg = "Could not find an appropriate base directory for saving data. Was the workspace initialised using set_up_workspace?"
        logger.error(msg)
        raise RuntimeError(msg)

    # check whether the table already exists and get the path if so
    existing_path = get_existing_table_save_dir(spark, catalog_db, table)

    if existing_path and mode in (IGNORE, ERROR, ERROR_IF_EXISTS):
        # TODO: check if data must be in the dir or whether the table just needs to be registered
        logger.warning(
            "Database table %s already exists and writer is set to %s mode, so no data would be written. Aborting.",
            catalog_db_table,
            mode,
        )
        return

    merge_or_overwrite_schema = "mergeSchema" if mode == APPEND else "overwriteSchema"
    # use to(schema) to ensure that the schema is saved with the dataframe
    writer = sdf.to(sdf.schema).write.format("delta").mode(mode).option(merge_or_overwrite_schema, "true")

    logger.info("Writing table %s in mode %s (rows=%d)", catalog_db_table, mode, sdf.count())
    logger.debug(sdf.printSchema())
    logger.debug(sdf.show(10, truncate=False))

    # use the standard spark saveAsTable, which will save to the appropriate subdirectory under the
    # designated namespace dir
    try:
        writer.saveAsTable(catalog_db_table)
        logger.info("Saved managed table %s (rows=%d)", catalog_db_table, sdf.count())
    except Exception:
        logger.exception("Error writing managed table %s", catalog_db_table)
        raise

def preview_or_skip(spark: SparkSession, catalog_db: str, table: str, limit: int = 20) -> None:
    """Preview table if it exists."""
    catalog_db_table = f"{catalog_db}.{table}"
    if not spark.catalog.tableExists(catalog_db_table):
        logger.info("Table %s not found. Skipping preview.", catalog_db_table)
        return

    logger.info("Preview for %s:", catalog_db_table)
    spark.sql(f"SELECT * FROM {catalog_db_table} LIMIT {limit}").show(truncate=False)