__generated_with = "0.13.15" # %% import sys import time from pyspark.sql.utils import AnalysisException sys.path.append('/opt/spark/work-dir/') from workflow_templates.spark.udf_manager import bootstrap_udfs from util import ( get_logger, observe_metrics, collect_metrics, log_info, log_error, forgiving_serializer, run_component, apply_data_quality, compute_dq_stats, enforce_error_threshold, build_dq_error_log, build_api_error_log, ERROR_LOG_SCHEMA, RetryConfig, with_retry, app_scoped_error_code, registry_error_code, rewrite_response_body_json_access, rewrite_response_body_json_access_if_json, ) from exception_utils import ( ErrorMessage, Severity, ConnectionException, AuthenticationException, SSLException, RateLimitException, ServiceUnavailableException, TimeoutException, ValidationException, SchemaMappingException, ExpressionException, MergeException, ConfigurationException, RetryExhaustedException, format_exception, mask_pii, mask_pii_dict, ) from py4j.protocol import Py4JJavaError from component_error_handler import handle_analysis_error, handle_java_error, classify_java_error from util import get_logger, observe_metrics, collect_metrics, log_info, log_error, forgiving_serializer, set_correlation_id, set_workflow_context from pyspark.sql.functions import udf from pyspark.sql.functions import count, expr, lit, input_file_name from pyspark.sql.types import StringType, IntegerType, MapType, StructType,StructField from postal.parser import parse_address import uuid from pathlib import Path from pyspark import SparkConf, Row from pyspark.sql import SparkSession from pyspark.sql.observation import Observation from pyspark import StorageLevel import os import pandas as pd import polars as pl import pyarrow as pa from pyspark.sql.functions import approx_count_distinct, avg, collect_list, collect_set, corr, count, countDistinct, covar_pop, covar_samp, first, kurtosis, last, max, mean, min, skewness, stddev, stddev_pop, stddev_samp, sum, var_pop, var_samp, variance,expr,to_json,struct, date_format, col, lit, when, regexp_replace, ltrim, lpad, format_number from functools import reduce from handle_structs_or_arrays import preprocess_then_expand import requests from requests.adapters import HTTPAdapter from urllib3.util.retry import Retry from jinja2 import Template import json import orjson from ocular_ai_sdk import OcularClient from ocular_ai_sdk.exceptions import ( OcularSDKException, AuthenticationError, ResourceNotFoundError ) from secrets_manager import SecretsManager from WorkflowManager import WorkflowDSL, WorkflowManager from KnowledgebaseManager import KnowledgebaseManager from gitea_client import GiteaClient, WorkspaceVersionedContent from FilesystemManager import FilesystemManager, SupportedFilesystemType from Materialization import Materialization init_start_time=time.time() LOGGER = get_logger() alias_str='abcdefghijklmnopqrstuvwxyz' workspace = os.getenv('WORKSPACE') or 'exp360uat' workflow = 'forecasting_workflow' execution_environment = os.getenv('EXECUTION_ENVIRONMENT') or 'CLUSTER' job_id = os.getenv("EXECUTION_ID") or str(uuid.uuid4()) retry_job_id = os.getenv("RETRY_EXECUTION_ID") or '' correlation_id = job_id set_correlation_id(correlation_id) set_workflow_context(workspace=workspace, workflow=workflow, job_id=job_id, retry_job_id=retry_job_id, execution_environment=execution_environment) log_info(LOGGER, f"Workspace: '{workspace}', Workflow: '{workflow}', Execution Environment: '{execution_environment}', Job Id: '{job_id}', Retry Job Id: '{retry_job_id}', Correlation Id: '{correlation_id}'") sm = SecretsManager(os.getenv('SECRET_MANAGER_URL'), os.getenv('SECRET_MANAGER_NAMESPACE'), os.getenv('SECRET_MANAGER_ENV'), os.getenv('SECRET_MANAGER_TOKEN')) secrets = sm.list_secrets(workspace) import dremio_operations dremio_operations.configure(secrets) import kb_query kb_query.configure(secrets) gitea_client=GiteaClient(os.getenv('GITEA_HOST'), os.getenv('GITEA_TOKEN'), os.getenv('GITEA_OWNER') or 'gitea_admin', os.getenv('GITEA_REPO') or 'tenant1') workspaceVersionedContent=WorkspaceVersionedContent(gitea_client) client = OcularClient( pat_token=secrets.get('OCULAR_AI_PAT_TOKEN') ) if 'AZURE_SERVICE_PRINCIPAL' in secrets: _storage_options=orjson.loads(secrets['AZURE_SERVICE_PRINCIPAL']) else: _storage_options = { 'key': secrets.get('S3_ACCESS_KEY'), 'secret': secrets.get('S3_SECRET_KEY'), 'region': secrets.get('S3_REGION') } filesystemManager = FilesystemManager.create(secrets.get('LAKEHOUSE_BUCKET'), storage_options=_storage_options) if retry_job_id: logs = Materialization.get_execution_history_by_job_id(filesystemManager, secrets.get('LAKEHOUSE_BUCKET'), workspace, workflow, retry_job_id, selected_components=['finalize']).to_dicts() if len(logs) == 1 and logs[0].get('metrics').get('execute_status') == 'SUCCESS': log_info(LOGGER, f"Workspace: '{workspace}', Workflow: '{workflow}', Execution Environment: '{execution_environment}', Job Id: '{job_id}' - Retry Job Id: '{retry_job_id}' was already successful. Hence exiting to forward processing to next in chain.") sys.exit(0) _conf = SparkConf() _params = { "spark.jars.ivy": "/opt/spark/.ivy2/", "spark.hadoop.fs.s3a.access.key": secrets.get('S3_ACCESS_KEY'), "spark.hadoop.fs.s3a.secret.key": secrets.get('S3_SECRET_KEY'), "spark.hadoop.fs.s3a.aws.region": secrets.get("S3_REGION") or "us-west-1", "spark.sql.catalog.dremio.warehouse" : secrets.get('LAKEHOUSE_BUCKET'), "spark.hadoop.fs.s3a.aws.credentials.provider": "com.amazonaws.auth.DefaultAWSCredentialsProviderChain", "spark.hadoop.fs.s3.aws.credentials.provider": "com.amazonaws.auth.DefaultAWSCredentialsProviderChain", "spark.sql.catalog.dremio" : "org.apache.iceberg.spark.SparkCatalog", "spark.sql.catalog.dremio.type" : "hadoop", "spark.hadoop.fs.s3a.impl": "org.apache.hadoop.fs.s3a.S3AFileSystem", "spark.hadoop.fs.s3.impl": "org.apache.hadoop.fs.s3a.S3AFileSystem", "spark.hadoop.fs.gs.impl": "com.google.cloud.hadoop.fs.gcs.GoogleHadoopFileSystem", "spark.sql.extensions": "org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions" } if filesystemManager.storage_type == SupportedFilesystemType.AZUREBLOB: _params[f"fs.azure.account.auth.type.{_storage_options['account_name']}.dfs.core.windows.net"] = "OAuth" _params[f"fs.azure.account.oauth.provider.type.{_storage_options['account_name']}.dfs.core.windows.net"] = "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider" _params[f"fs.azure.account.oauth2.client.id.{_storage_options['account_name']}.dfs.core.windows.net"] = _storage_options['client_id'] _params[f"fs.azure.account.oauth2.client.secret.{_storage_options['account_name']}.dfs.core.windows.net"] = _storage_options['client_secret'] _params[f"fs.azure.account.oauth2.client.endpoint.{_storage_options['account_name']}.dfs.core.windows.net"] = f"https://login.microsoftonline.com/{_storage_options['tenant_id']}/oauth2/v2.0/token" _conf.setAll(list(_params.items())) spark = SparkSession.builder.appName(workspace).config(conf=_conf).getOrCreate() bootstrap_udfs(spark) materialization = Materialization(spark, secrets.get('LAKEHOUSE_BUCKET'), workspace, workflow, job_id, retry_job_id, execution_environment, LOGGER) init_dependency_key="init" init_end_time=time.time() # %% readCustomers_start_time=time.time() readCustomers_fail_on_error="" try: _readCustomers_options = { 'jdbc':{ 'dbtable': """customer""", 'url':secrets.get(''), 'driver':'' }, 'kafka' : { 'kafka.bootstrap.servers' : secrets.get('OCULAR_KAFKA_BOOTSTRAP_SERVERS'), 'subscribe' : '', 'startingOffsets' : 'earliest' }, 'cobol' : { 'copybook' : '', 'encoding' : '', 'is_text': False, 'schema_retention_policy' : 'collapse_root' } } _reader = spark.read.format('iceberg') _readCustomers_load_path = 'dremio.customer' _readCustomers_input_data = { "component": "readCustomers", "format": "iceberg", "iceberg_catalog": """dremio""", "table_name": """customer""", } try: readCustomers_df = _reader.load(_readCustomers_load_path) readCustomers_df = readCustomers_df.withColumn("readCustomers_input_file", input_file_name()) # Force partition evaluation to surface lazy errors (e.g. glob matches 0 files) readCustomers_df.rdd.getNumPartitions() except AnalysisException as e: handle_analysis_error( e, component_name="readCustomers", message=f"Failed to load source 'readCustomers' ({_readCustomers_load_path}): {e!s}", job_id=job_id, workspace=workspace, workflow=workflow, execution_environment=execution_environment, extra_details={"format": "iceberg", "load_path": _readCustomers_load_path}, input_data=_readCustomers_input_data, ) except Py4JJavaError as e: handle_java_error( e, component_name="readCustomers", operation="load", format_name="iceberg", path=_readCustomers_load_path, job_id=job_id, workspace=workspace, workflow=workflow, execution_environment=execution_environment, input_data=_readCustomers_input_data, ) readCustomers_df, readCustomers_observer = observe_metrics("readCustomers_df", readCustomers_df) readCustomers_df.createOrReplaceTempView('readCustomers_df') readCustomers_dependency_key="readCustomers" readCustomers_execute_status="SUCCESS" except Exception as e: readCustomers_error = e log_error(LOGGER, f"Component readCustomers Failed", e, component_name="readCustomers") readCustomers_execute_status="ERROR" raise e finally: readCustomers_end_time=time.time() # %% filterActiveCustomers_start_time=time.time() print(readCustomers_df.columns) filterActiveCustomers_fail_on_error="True" try: _filterActiveCustomers_condition = rewrite_response_body_json_access_if_json(readCustomers_df, """TRIM(UPPER(status)) = \'ACTIVE\'""") filterActiveCustomers_df = spark.sql(f"select * from readCustomers_df where {_filterActiveCustomers_condition}") filterActiveCustomers_df, filterActiveCustomers_observer = observe_metrics("filterActiveCustomers_df", filterActiveCustomers_df) filterActiveCustomers_df.createOrReplaceTempView('filterActiveCustomers_df') filterActiveCustomers_dependency_key="filterActiveCustomers" print(readCustomers_dependency_key) filterActiveCustomers_execute_status="SUCCESS" except Exception as e: filterActiveCustomers_error = e log_error(LOGGER, f"Component filterActiveCustomers Failed", e, component_name="filterActiveCustomers") filterActiveCustomers_execute_status="ERROR" raise e finally: filterActiveCustomers_end_time=time.time() # %% getBills_start_time=time.time() getBills_fail_on_error="True" getBills_observer = Observation("getBills_df") try: _getBills_url = 'https://fw-gateway:8200/fw-notification/outbound-message-config/publish' _getBills_headers = dict() for _getBills_k,_getBills_v in {'Content-Type': {'value': 'application/json', 'secret': None}, 'Content-type': {'value': 'application/json', 'secret': None}, 'api-key': {'value': None, 'secret': 'OCULAR_API_KEY'}, 'x-tenantCode': {'value': 'UTILITIES', 'secret': None}}.items() : if(_getBills_v.get('value') is not None and _getBills_v.get('value') != ''): _getBills_headers[_getBills_k] = _getBills_v.get('value') elif(_getBills_v.get('secret') is not None and _getBills_v.get('secret') != ''): _getBills_headers[_getBills_k] = secrets.get(_getBills_v.get('secret')) _timeout=(5, 90) _getBills_out_schema = StructType( list(filterActiveCustomers_df.schema.fields) + [ StructField("response_body", StringType(), True), StructField("request_body", StringType(), True), StructField("response_status_code", IntegerType(), True), StructField("api_error", StringType(), True), StructField("api_error_code", StringType(), True), ] ) def _getBills_call_api(row, _session): body_dict = row.asDict(recursive=True) template = Template('''{ "outMsgConfigCode": "EXP_ACCOUNT_BILL_HISTORY", "msgData": { "accountId": "{{account_id}}", "numberOfMonthPast": "24" } }''') _request_body_json = None try: body = json.loads(template.render(**body_dict)) print("request : "+ json.dumps(body)) _request_body_json = json.dumps(body, default=str) _retry_cfg = RetryConfig(retries=2, backoff_seconds=1.0, backoff_multiplier=2.0, max_backoff_seconds=10.0) def _do_request(): return _session.post(_getBills_url, headers=_getBills_headers, json=body, params={}, verify=False, timeout=_timeout) response = with_retry(_do_request, _retry_cfg) _status = int(response.status_code) try: data = response.json() _response_body_json = json.dumps(data, default=str) except Exception: _response_body_json = response.text if _status >= 400: _api_ctx = { "job_id": job_id, "workspace": workspace, "workflow": workflow, "execution_environment": execution_environment, "url": _getBills_url, "method": "POST", "http_status": _status, } if _status in (401, 403, 440): _typed_exc = AuthenticationException( message=f"HTTP {_status} from getBills: authentication/authorization failed", error_code="NET-AUTH-001", source="getBills", correlation_id=job_id, severity=Severity.ERROR, details=_api_ctx, mask_pii=True, ) elif _status == 429: _typed_exc = RateLimitException( message=f"HTTP 429 from getBills: rate limit exceeded", error_code="NET-RATE-001", source="getBills", correlation_id=job_id, severity=Severity.WARNING, details=_api_ctx, mask_pii=True, ) elif _status == 503: _typed_exc = ServiceUnavailableException( message=f"HTTP 503 from getBills: service unavailable", error_code="NET-SVC-001", source="getBills", correlation_id=job_id, severity=Severity.ERROR, details=_api_ctx, mask_pii=True, ) else: _typed_exc = ConnectionException( message=f"HTTP {_status} from getBills", error_code="NET-CON-001", source="getBills", correlation_id=job_id, severity=Severity.ERROR, details=_api_ctx, mask_pii=True, ) _scoped_code = app_scoped_error_code(getattr(_typed_exc, 'error_code', 'SYS-UNK-001')) _err_payload = { "error_code": _scoped_code, "message": getattr(_typed_exc, 'message', str(_typed_exc)), "response": _response_body_json, } merged = { **body_dict, "response_body": _response_body_json, "request_body": _request_body_json, "response_status_code": _status, "api_error": json.dumps(_err_payload, default=str), "api_error_code": _scoped_code, } return json.dumps(merged, default=str) print("response : " + _response_body_json) merged = { **body_dict, "response_body": _response_body_json, "request_body": _request_body_json, "response_status_code": _status, "api_error": None, "api_error_code": None, } return json.dumps(merged, default=str) except requests.exceptions.Timeout as e: _typed_exc = TimeoutException( message=f"Request timeout for getBills: {e!s}", error_code="RES-TMO-001", source="getBills", correlation_id=job_id, severity=Severity.ERROR, details={ "job_id": job_id, "workspace": workspace, "workflow": workflow, "url": _getBills_url, "method": "POST", }, cause=e, mask_pii=True, ) _scoped_code = app_scoped_error_code(getattr(_typed_exc, 'error_code', 'SYS-UNK-001')) _err_payload = { "error_code": _scoped_code, "message": getattr(_typed_exc, 'message', str(_typed_exc)), "response": None, } merged = { **body_dict, "response_body": None, "request_body": _request_body_json, "response_status_code": None, "api_error": json.dumps(_err_payload, default=str), "api_error_code": _scoped_code, } return json.dumps(merged, default=str) except requests.exceptions.SSLError as e: _typed_exc = SSLException( message=f"SSL error for getBills: {e!s}", error_code="NET-SSL-001", source="getBills", correlation_id=job_id, severity=Severity.ERROR, details={ "job_id": job_id, "workspace": workspace, "workflow": workflow, "url": _getBills_url, "method": "POST", }, cause=e, mask_pii=True, ) _scoped_code = app_scoped_error_code(getattr(_typed_exc, 'error_code', 'SYS-UNK-001')) _err_payload = { "error_code": _scoped_code, "message": getattr(_typed_exc, 'message', str(_typed_exc)), "response": None, } merged = { **body_dict, "response_body": None, "request_body": _request_body_json, "response_status_code": None, "api_error": json.dumps(_err_payload, default=str), "api_error_code": _scoped_code, } return json.dumps(merged, default=str) except requests.exceptions.ConnectionError as e: _typed_exc = ConnectionException( message=f"Connection error for getBills: {e!s}", error_code="NET-CON-001", source="getBills", correlation_id=job_id, severity=Severity.ERROR, details={ "job_id": job_id, "workspace": workspace, "workflow": workflow, "url": _getBills_url, "method": "POST", }, cause=e, mask_pii=True, ) _scoped_code = app_scoped_error_code(getattr(_typed_exc, 'error_code', 'SYS-UNK-001')) _err_payload = { "error_code": _scoped_code, "message": getattr(_typed_exc, 'message', str(_typed_exc)), "response": None, } merged = { **body_dict, "response_body": None, "request_body": _request_body_json, "response_status_code": None, "api_error": json.dumps(_err_payload, default=str), "api_error_code": _scoped_code, } return json.dumps(merged, default=str) except RetryExhaustedException as e: _typed_exc = e _scoped_code = app_scoped_error_code(getattr(_typed_exc, 'error_code', 'SYS-RET-001')) _err_payload = { "error_code": _scoped_code, "message": getattr(_typed_exc, 'message', str(_typed_exc)), "response": None, } merged = { **body_dict, "response_body": None, "request_body": _request_body_json, "response_status_code": None, "api_error": json.dumps(_err_payload, default=str), "api_error_code": _scoped_code, } return json.dumps(merged, default=str) except Exception as e: import traceback as _tb _err_trace = str(e) + "\n" + _tb.format_exc() _scoped_code = app_scoped_error_code("SYS-UNK-001") _err_payload = { "error_code": _scoped_code, "message": _err_trace, "response": None, } merged = { **body_dict, "response_body": None, "request_body": _request_body_json, "response_status_code": None, "api_error": json.dumps(_err_payload, default=str), "api_error_code": _scoped_code, } return json.dumps(merged, default=str) def _getBills_partition(rows, _call_api=_getBills_call_api): import requests as _req from requests.adapters import HTTPAdapter as _HA from urllib3.util.retry import Retry as _Retry _rs = _Retry(total=3, connect=2, read=2, backoff_factor=1, respect_retry_after_header=True) _ad = _HA(max_retries=_rs) _session = _req.Session() _session.mount("https://", _ad) _session.mount("http://", _ad) for row in rows: yield _call_api(row, _session) _rdd=filterActiveCustomers_df.rdd.mapPartitions(_getBills_partition).persist() try: getBills_df = spark.read.schema(_getBills_out_schema).json(_rdd) except Py4JJavaError as e: _java_msg = str(e.java_exception) if hasattr(e, 'java_exception') else str(e) raise ConfigurationException( message=f"Failed to parse API response RDD for 'getBills': {_java_msg}", error_code="SYS-CFG-001", source="getBills", correlation_id=job_id, severity=Severity.ERROR, details={ "job_id": job_id, "workspace": workspace, "workflow": workflow, "execution_environment": execution_environment, "java_exception": _java_msg[:2000], }, cause=e, mask_pii=True, ) from e getBills_df.persist() getBills_df.count() # Force API execution (Spark action) getBills_df, getBills_observer = observe_metrics("getBills_df", getBills_df, getBills_observer) getBills_df.createOrReplaceTempView('getBills_df') _getBills_error_log = build_api_error_log( getBills_df, job_execution_id=job_id, step_name="getBills", workspace=workspace, workflow=workflow, execution_environment=execution_environment, retry_job_id=retry_job_id, http_method="POST", ) _getBills_error_log.createOrReplaceTempView("_getBills_error_temp") _getBills_error_log.writeTo("dremio.error_dlq").append() getBills_dependency_key="getBills" print(filterActiveCustomers_dependency_key) getBills_execute_status="SUCCESS" except Exception as e: getBills_error = e log_error(LOGGER, f"Component getBills Failed", e, component_name="getBills") getBills_execute_status="ERROR" raise e finally: getBills_end_time=time.time() # %% readLatestBillIds_start_time=time.time() readLatestBillIds_fail_on_error="" try: _readLatestBillIds_options = { 'jdbc':{ 'dbtable': """bills""", 'url':secrets.get(''), 'driver':'' }, 'kafka' : { 'kafka.bootstrap.servers' : secrets.get('OCULAR_KAFKA_BOOTSTRAP_SERVERS'), 'subscribe' : '', 'startingOffsets' : 'earliest' }, 'cobol' : { 'copybook' : '', 'encoding' : '', 'is_text': False, 'schema_retention_policy' : 'collapse_root' } } _reader = spark.read.format('iceberg') _readLatestBillIds_load_path = 'dremio.bills' _readLatestBillIds_input_data = { "component": "readLatestBillIds", "format": "iceberg", "iceberg_catalog": """dremio""", "table_name": """bills""", } try: readLatestBillIds_df = _reader.load(_readLatestBillIds_load_path) readLatestBillIds_df = readLatestBillIds_df.withColumn("readLatestBillIds_input_file", input_file_name()) # Force partition evaluation to surface lazy errors (e.g. glob matches 0 files) readLatestBillIds_df.rdd.getNumPartitions() except AnalysisException as e: handle_analysis_error( e, component_name="readLatestBillIds", message=f"Failed to load source 'readLatestBillIds' ({_readLatestBillIds_load_path}): {e!s}", job_id=job_id, workspace=workspace, workflow=workflow, execution_environment=execution_environment, extra_details={"format": "iceberg", "load_path": _readLatestBillIds_load_path}, input_data=_readLatestBillIds_input_data, ) except Py4JJavaError as e: handle_java_error( e, component_name="readLatestBillIds", operation="load", format_name="iceberg", path=_readLatestBillIds_load_path, job_id=job_id, workspace=workspace, workflow=workflow, execution_environment=execution_environment, input_data=_readLatestBillIds_input_data, ) readLatestBillIds_df, readLatestBillIds_observer = observe_metrics("readLatestBillIds_df", readLatestBillIds_df) readLatestBillIds_df.createOrReplaceTempView('readLatestBillIds_df') readLatestBillIds_dependency_key="readLatestBillIds" readLatestBillIds_execute_status="SUCCESS" except Exception as e: readLatestBillIds_error = e log_error(LOGGER, f"Component readLatestBillIds Failed", e, component_name="readLatestBillIds") readLatestBillIds_execute_status="ERROR" raise e finally: readLatestBillIds_end_time=time.time() # %% MapLatestBill_start_time=time.time() MapLatestBill_fail_on_error="True" try: _MapLatestBill_select_clause=[] _MapLatestBill_expr = """from_json( get_json_object(response_body, \'$.data\'), \'struct< accountId:string, numberOfMonthPast:string, output:struct< bills:array> > >\' )""".replace("input_file_name()", "input_file") _MapLatestBill_expr = _MapLatestBill_expr.replace("_dq_source_file", "input_file") if "." in _MapLatestBill_expr: _MapLatestBill_expr = rewrite_response_body_json_access(_MapLatestBill_expr) _MapLatestBill_select_clause.append(f"{_MapLatestBill_expr} AS accounts") _MapLatestBill_expr = """account_id""".replace("input_file_name()", "input_file") _MapLatestBill_expr = _MapLatestBill_expr.replace("_dq_source_file", "input_file") if "." in _MapLatestBill_expr: _MapLatestBill_expr = rewrite_response_body_json_access(_MapLatestBill_expr) _MapLatestBill_select_clause.append(f"{_MapLatestBill_expr} AS account_id") _MapLatestBill_mapping_sql = ("SELECT " + ', '.join(_MapLatestBill_select_clause) + " FROM getBills_df").replace("{job_id}", f"'{job_id}'") _MapLatestBill_input_data = { "component": "MapLatestBill", "datasource": "getBills", "include_existing_columns": False, "to_schema_field_count": 2, } try: MapLatestBill_df = spark.sql(_MapLatestBill_mapping_sql) except AnalysisException as e: handle_analysis_error( e, component_name="MapLatestBill", error_code="TRF-MAP-002", exception_class=SchemaMappingException, message=f"Spark analysis error during MapLatestBill mapping: {e!s}", job_id=job_id, workspace=workspace, workflow=workflow, execution_environment=execution_environment, extra_details={"retry_job_id": retry_job_id or None, "sql_preview": _MapLatestBill_mapping_sql[:2000]}, input_data=_MapLatestBill_input_data, ) except Py4JJavaError as e: handle_java_error( e, component_name="MapLatestBill", operation="mapping SQL", format_name="sql", job_id=job_id, workspace=workspace, workflow=workflow, execution_environment=execution_environment, override_class=ExpressionException, override_code="TRF-EXP-001", extra_details={"retry_job_id": retry_job_id or None, "sql_preview": _MapLatestBill_mapping_sql[:2000]}, input_data=_MapLatestBill_input_data, ) MapLatestBill_df, MapLatestBill_observer = observe_metrics("MapLatestBill_df", MapLatestBill_df) MapLatestBill_df.createOrReplaceTempView("MapLatestBill_df") MapLatestBill_dependency_key="MapLatestBill" print(getBills_dependency_key) MapLatestBill_execute_status="SUCCESS" except Exception as e: MapLatestBill_error = e log_error(LOGGER, f"Component MapLatestBill Failed", e, component_name="MapLatestBill") MapLatestBill_execute_status="ERROR" raise e finally: MapLatestBill_end_time=time.time() # %% data_mapper__3_start_time=time.time() data_mapper__3_fail_on_error="True" try: _data_mapper__3_select_clause=[] _data_mapper__3_expr = """account_id""".replace("input_file_name()", "input_file") _data_mapper__3_expr = _data_mapper__3_expr.replace("_dq_source_file", "input_file") if "." in _data_mapper__3_expr: _data_mapper__3_expr = rewrite_response_body_json_access(_data_mapper__3_expr) _data_mapper__3_select_clause.append(f"{_data_mapper__3_expr} AS account_id") _data_mapper__3_expr = """accounts.output.bills""".replace("input_file_name()", "input_file") _data_mapper__3_expr = _data_mapper__3_expr.replace("_dq_source_file", "input_file") if "." in _data_mapper__3_expr: _data_mapper__3_expr = rewrite_response_body_json_access(_data_mapper__3_expr) _data_mapper__3_select_clause.append(f"{_data_mapper__3_expr} AS bills") _data_mapper__3_mapping_sql = ("SELECT " + ', '.join(_data_mapper__3_select_clause) + " FROM MapLatestBill_df").replace("{job_id}", f"'{job_id}'") _data_mapper__3_input_data = { "component": "data_mapper__3", "datasource": "MapLatestBill", "include_existing_columns": False, "to_schema_field_count": 2, } try: data_mapper__3_df = spark.sql(_data_mapper__3_mapping_sql) except AnalysisException as e: handle_analysis_error( e, component_name="data_mapper__3", error_code="TRF-MAP-002", exception_class=SchemaMappingException, message=f"Spark analysis error during data_mapper__3 mapping: {e!s}", job_id=job_id, workspace=workspace, workflow=workflow, execution_environment=execution_environment, extra_details={"retry_job_id": retry_job_id or None, "sql_preview": _data_mapper__3_mapping_sql[:2000]}, input_data=_data_mapper__3_input_data, ) except Py4JJavaError as e: handle_java_error( e, component_name="data_mapper__3", operation="mapping SQL", format_name="sql", job_id=job_id, workspace=workspace, workflow=workflow, execution_environment=execution_environment, override_class=ExpressionException, override_code="TRF-EXP-001", extra_details={"retry_job_id": retry_job_id or None, "sql_preview": _data_mapper__3_mapping_sql[:2000]}, input_data=_data_mapper__3_input_data, ) data_mapper__3_df, data_mapper__3_observer = observe_metrics("data_mapper__3_df", data_mapper__3_df) data_mapper__3_df.createOrReplaceTempView("data_mapper__3_df") data_mapper__3_dependency_key="data_mapper__3" print(MapLatestBill_dependency_key) data_mapper__3_execute_status="SUCCESS" except Exception as e: data_mapper__3_error = e log_error(LOGGER, f"Component data_mapper__3 Failed", e, component_name="data_mapper__3") data_mapper__3_execute_status="ERROR" raise e finally: data_mapper__3_end_time=time.time() # %% data_mapper__4_start_time=time.time() data_mapper__4_fail_on_error="True" try: _data_mapper__4_select_clause=[] _data_mapper__4_expr = """account_id""".replace("input_file_name()", "input_file") _data_mapper__4_expr = _data_mapper__4_expr.replace("_dq_source_file", "input_file") if "." in _data_mapper__4_expr: _data_mapper__4_expr = rewrite_response_body_json_access(_data_mapper__4_expr) _data_mapper__4_select_clause.append(f"{_data_mapper__4_expr} AS account_id") _data_mapper__4_expr = """explode(bills)""".replace("input_file_name()", "input_file") _data_mapper__4_expr = _data_mapper__4_expr.replace("_dq_source_file", "input_file") if "." in _data_mapper__4_expr: _data_mapper__4_expr = rewrite_response_body_json_access(_data_mapper__4_expr) _data_mapper__4_select_clause.append(f"{_data_mapper__4_expr} AS final_bills") _data_mapper__4_mapping_sql = ("SELECT " + ', '.join(_data_mapper__4_select_clause) + " FROM data_mapper__3_df").replace("{job_id}", f"'{job_id}'") _data_mapper__4_input_data = { "component": "data_mapper__4", "datasource": "data_mapper__3", "include_existing_columns": False, "to_schema_field_count": 2, } try: data_mapper__4_df = spark.sql(_data_mapper__4_mapping_sql) except AnalysisException as e: handle_analysis_error( e, component_name="data_mapper__4", error_code="TRF-MAP-002", exception_class=SchemaMappingException, message=f"Spark analysis error during data_mapper__4 mapping: {e!s}", job_id=job_id, workspace=workspace, workflow=workflow, execution_environment=execution_environment, extra_details={"retry_job_id": retry_job_id or None, "sql_preview": _data_mapper__4_mapping_sql[:2000]}, input_data=_data_mapper__4_input_data, ) except Py4JJavaError as e: handle_java_error( e, component_name="data_mapper__4", operation="mapping SQL", format_name="sql", job_id=job_id, workspace=workspace, workflow=workflow, execution_environment=execution_environment, override_class=ExpressionException, override_code="TRF-EXP-001", extra_details={"retry_job_id": retry_job_id or None, "sql_preview": _data_mapper__4_mapping_sql[:2000]}, input_data=_data_mapper__4_input_data, ) data_mapper__4_df, data_mapper__4_observer = observe_metrics("data_mapper__4_df", data_mapper__4_df) data_mapper__4_df.createOrReplaceTempView("data_mapper__4_df") data_mapper__4_dependency_key="data_mapper__4" print(data_mapper__3_dependency_key) data_mapper__4_execute_status="SUCCESS" except Exception as e: data_mapper__4_error = e log_error(LOGGER, f"Component data_mapper__4 Failed", e, component_name="data_mapper__4") data_mapper__4_execute_status="ERROR" raise e finally: data_mapper__4_end_time=time.time() # %% data_mapper__5_start_time=time.time() data_mapper__5_fail_on_error="True" try: _data_mapper__5_select_clause=[] _data_mapper__5_expr = """account_id""".replace("input_file_name()", "input_file") _data_mapper__5_expr = _data_mapper__5_expr.replace("_dq_source_file", "input_file") if "." in _data_mapper__5_expr: _data_mapper__5_expr = rewrite_response_body_json_access(_data_mapper__5_expr) _data_mapper__5_select_clause.append(f"{_data_mapper__5_expr} AS account_id") _data_mapper__5_expr = """current_timestamp()""".replace("input_file_name()", "input_file") _data_mapper__5_expr = _data_mapper__5_expr.replace("_dq_source_file", "input_file") if "." in _data_mapper__5_expr: _data_mapper__5_expr = rewrite_response_body_json_access(_data_mapper__5_expr) _data_mapper__5_select_clause.append(f"{_data_mapper__5_expr} AS created_at") _data_mapper__5_expr = """uuid()""".replace("input_file_name()", "input_file") _data_mapper__5_expr = _data_mapper__5_expr.replace("_dq_source_file", "input_file") if "." in _data_mapper__5_expr: _data_mapper__5_expr = rewrite_response_body_json_access(_data_mapper__5_expr) _data_mapper__5_select_clause.append(f"{_data_mapper__5_expr} AS id") _data_mapper__5_expr = """to_date(final_bills.billDate)""".replace("input_file_name()", "input_file") _data_mapper__5_expr = _data_mapper__5_expr.replace("_dq_source_file", "input_file") if "." in _data_mapper__5_expr: _data_mapper__5_expr = rewrite_response_body_json_access(_data_mapper__5_expr) _data_mapper__5_select_clause.append(f"{_data_mapper__5_expr} AS bill_date") _data_mapper__5_expr = """to_date(final_bills.dueDate)""".replace("input_file_name()", "input_file") _data_mapper__5_expr = _data_mapper__5_expr.replace("_dq_source_file", "input_file") if "." in _data_mapper__5_expr: _data_mapper__5_expr = rewrite_response_body_json_access(_data_mapper__5_expr) _data_mapper__5_select_clause.append(f"{_data_mapper__5_expr} AS due_date") _data_mapper__5_expr = """final_bills.billStatus""".replace("input_file_name()", "input_file") _data_mapper__5_expr = _data_mapper__5_expr.replace("_dq_source_file", "input_file") if "." in _data_mapper__5_expr: _data_mapper__5_expr = rewrite_response_body_json_access(_data_mapper__5_expr) _data_mapper__5_select_clause.append(f"{_data_mapper__5_expr} AS bill_status") _data_mapper__5_expr = """final_bills.billId""".replace("input_file_name()", "input_file") _data_mapper__5_expr = _data_mapper__5_expr.replace("_dq_source_file", "input_file") if "." in _data_mapper__5_expr: _data_mapper__5_expr = rewrite_response_body_json_access(_data_mapper__5_expr) _data_mapper__5_select_clause.append(f"{_data_mapper__5_expr} AS mapper_bill_id") _data_mapper__5_expr = """cast(replace(replace(final_bills.amount, \'$\', \'\'), \',\', \'\')AS decimal(10, 2))""".replace("input_file_name()", "input_file") _data_mapper__5_expr = _data_mapper__5_expr.replace("_dq_source_file", "input_file") if "." in _data_mapper__5_expr: _data_mapper__5_expr = rewrite_response_body_json_access(_data_mapper__5_expr) _data_mapper__5_select_clause.append(f"{_data_mapper__5_expr} AS amount_value") _data_mapper__5_expr = """final_bills.amount""".replace("input_file_name()", "input_file") _data_mapper__5_expr = _data_mapper__5_expr.replace("_dq_source_file", "input_file") if "." in _data_mapper__5_expr: _data_mapper__5_expr = rewrite_response_body_json_access(_data_mapper__5_expr) _data_mapper__5_select_clause.append(f"{_data_mapper__5_expr} AS amount") _data_mapper__5_expr = """final_bills.billStatusName""".replace("input_file_name()", "input_file") _data_mapper__5_expr = _data_mapper__5_expr.replace("_dq_source_file", "input_file") if "." in _data_mapper__5_expr: _data_mapper__5_expr = rewrite_response_body_json_access(_data_mapper__5_expr) _data_mapper__5_select_clause.append(f"{_data_mapper__5_expr} AS bill_status_name") _data_mapper__5_expr = """final_bills.completionDttm""".replace("input_file_name()", "input_file") _data_mapper__5_expr = _data_mapper__5_expr.replace("_dq_source_file", "input_file") if "." in _data_mapper__5_expr: _data_mapper__5_expr = rewrite_response_body_json_access(_data_mapper__5_expr) _data_mapper__5_select_clause.append(f"{_data_mapper__5_expr} AS completion_dttm") _data_mapper__5_mapping_sql = ("SELECT " + ', '.join(_data_mapper__5_select_clause) + " FROM data_mapper__4_df").replace("{job_id}", f"'{job_id}'") _data_mapper__5_input_data = { "component": "data_mapper__5", "datasource": "data_mapper__4", "include_existing_columns": False, "to_schema_field_count": 11, } try: data_mapper__5_df = spark.sql(_data_mapper__5_mapping_sql) except AnalysisException as e: handle_analysis_error( e, component_name="data_mapper__5", error_code="TRF-MAP-002", exception_class=SchemaMappingException, message=f"Spark analysis error during data_mapper__5 mapping: {e!s}", job_id=job_id, workspace=workspace, workflow=workflow, execution_environment=execution_environment, extra_details={"retry_job_id": retry_job_id or None, "sql_preview": _data_mapper__5_mapping_sql[:2000]}, input_data=_data_mapper__5_input_data, ) except Py4JJavaError as e: handle_java_error( e, component_name="data_mapper__5", operation="mapping SQL", format_name="sql", job_id=job_id, workspace=workspace, workflow=workflow, execution_environment=execution_environment, override_class=ExpressionException, override_code="TRF-EXP-001", extra_details={"retry_job_id": retry_job_id or None, "sql_preview": _data_mapper__5_mapping_sql[:2000]}, input_data=_data_mapper__5_input_data, ) data_mapper__5_df, data_mapper__5_observer = observe_metrics("data_mapper__5_df", data_mapper__5_df) data_mapper__5_df.createOrReplaceTempView("data_mapper__5_df") data_mapper__5_dependency_key="data_mapper__5" print(data_mapper__4_dependency_key) data_mapper__5_execute_status="SUCCESS" except Exception as e: data_mapper__5_error = e log_error(LOGGER, f"Component data_mapper__5 Failed", e, component_name="data_mapper__5") data_mapper__5_execute_status="ERROR" raise e finally: data_mapper__5_end_time=time.time() # %% data_join__0_start_time=time.time() data_join__0_fail_on_error="" try: _data_join__0_select_clause, _data_join__0_from_clause = list(map(lambda i: (i, 'a'), data_mapper__5_df.columns)), ['data_mapper__5_df a'] _data_join__0_select_clause.extend(list(map(lambda i: (i, alias_str[1]), readLatestBillIds_df.columns))) _data_join__0_from_clause.append(WorkflowManager.build_join_clause({'with': 'readLatestBillIds', 'joinColumns': [{'account_Id': 'account_Id'}, {'mapper_bill_id': 'bill_id'}], 'how': 'left outer'}, alias_str[0], alias_str[1])) _data_join__0_from_clause_str=''.join(_data_join__0_from_clause) _data_join__0_select_clause_str=', '.join(map(lambda i: f"{i[1]}.`{i[0]}`", reversed(dict(reversed(_data_join__0_select_clause)).items()))) data_join__0_df=spark.sql("SELECT " + _data_join__0_select_clause_str + " FROM " + _data_join__0_from_clause_str) data_join__0_df, data_join__0_observer = observe_metrics("data_join__0_df", data_join__0_df) data_join__0_df.createOrReplaceTempView("data_join__0_df") data_join__0_dependency_key="data_join__0" print(readLatestBillIds_dependency_key) print(data_mapper__5_dependency_key) data_join__0_execute_status="SUCCESS" except Exception as e: data_join__0_error = e log_error(LOGGER, f"Component data_join__0 Failed", e, component_name="data_join__0") data_join__0_execute_status="ERROR" raise e finally: data_join__0_end_time=time.time() # %% filter__1_start_time=time.time() print(data_join__0_df.columns) filter__1_fail_on_error="True" try: _filter__1_condition = rewrite_response_body_json_access_if_json(data_join__0_df, """bill_id IS NULL OR mapper_bill_id <> bill_id""") filter__1_df = spark.sql(f"select * from data_join__0_df where {_filter__1_condition}") filter__1_df, filter__1_observer = observe_metrics("filter__1_df", filter__1_df) filter__1_df.createOrReplaceTempView('filter__1_df') filter__1_dependency_key="filter__1" print(data_join__0_dependency_key) filter__1_execute_status="SUCCESS" except Exception as e: filter__1_error = e log_error(LOGGER, f"Component filter__1 Failed", e, component_name="filter__1") filter__1_execute_status="ERROR" raise e finally: filter__1_end_time=time.time() # %% BillWriterMapper_start_time=time.time() BillWriterMapper_fail_on_error="True" try: _BillWriterMapper_select_clause=[] _BillWriterMapper_expr = """account_id""".replace("input_file_name()", "input_file") _BillWriterMapper_expr = _BillWriterMapper_expr.replace("_dq_source_file", "input_file") if "." in _BillWriterMapper_expr: _BillWriterMapper_expr = rewrite_response_body_json_access(_BillWriterMapper_expr) _BillWriterMapper_select_clause.append(f"{_BillWriterMapper_expr} AS account_id") _BillWriterMapper_expr = """mapper_bill_id""".replace("input_file_name()", "input_file") _BillWriterMapper_expr = _BillWriterMapper_expr.replace("_dq_source_file", "input_file") if "." in _BillWriterMapper_expr: _BillWriterMapper_expr = rewrite_response_body_json_access(_BillWriterMapper_expr) _BillWriterMapper_select_clause.append(f"{_BillWriterMapper_expr} AS bill_id") _BillWriterMapper_expr = """bill_date""".replace("input_file_name()", "input_file") _BillWriterMapper_expr = _BillWriterMapper_expr.replace("_dq_source_file", "input_file") if "." in _BillWriterMapper_expr: _BillWriterMapper_expr = rewrite_response_body_json_access(_BillWriterMapper_expr) _BillWriterMapper_select_clause.append(f"{_BillWriterMapper_expr} AS bill_date") _BillWriterMapper_expr = """bill_status""".replace("input_file_name()", "input_file") _BillWriterMapper_expr = _BillWriterMapper_expr.replace("_dq_source_file", "input_file") if "." in _BillWriterMapper_expr: _BillWriterMapper_expr = rewrite_response_body_json_access(_BillWriterMapper_expr) _BillWriterMapper_select_clause.append(f"{_BillWriterMapper_expr} AS bill_status") _BillWriterMapper_expr = """due_date""".replace("input_file_name()", "input_file") _BillWriterMapper_expr = _BillWriterMapper_expr.replace("_dq_source_file", "input_file") if "." in _BillWriterMapper_expr: _BillWriterMapper_expr = rewrite_response_body_json_access(_BillWriterMapper_expr) _BillWriterMapper_select_clause.append(f"{_BillWriterMapper_expr} AS due_date") _BillWriterMapper_expr = """current_timestamp()""".replace("input_file_name()", "input_file") _BillWriterMapper_expr = _BillWriterMapper_expr.replace("_dq_source_file", "input_file") if "." in _BillWriterMapper_expr: _BillWriterMapper_expr = rewrite_response_body_json_access(_BillWriterMapper_expr) _BillWriterMapper_select_clause.append(f"{_BillWriterMapper_expr} AS created_at") _BillWriterMapper_expr = """id""".replace("input_file_name()", "input_file") _BillWriterMapper_expr = _BillWriterMapper_expr.replace("_dq_source_file", "input_file") if "." in _BillWriterMapper_expr: _BillWriterMapper_expr = rewrite_response_body_json_access(_BillWriterMapper_expr) _BillWriterMapper_select_clause.append(f"{_BillWriterMapper_expr} AS id") _BillWriterMapper_expr = """amount""".replace("input_file_name()", "input_file") _BillWriterMapper_expr = _BillWriterMapper_expr.replace("_dq_source_file", "input_file") if "." in _BillWriterMapper_expr: _BillWriterMapper_expr = rewrite_response_body_json_access(_BillWriterMapper_expr) _BillWriterMapper_select_clause.append(f"{_BillWriterMapper_expr} AS amount") _BillWriterMapper_expr = """amount_value""".replace("input_file_name()", "input_file") _BillWriterMapper_expr = _BillWriterMapper_expr.replace("_dq_source_file", "input_file") if "." in _BillWriterMapper_expr: _BillWriterMapper_expr = rewrite_response_body_json_access(_BillWriterMapper_expr) _BillWriterMapper_select_clause.append(f"{_BillWriterMapper_expr} AS amount_value") _BillWriterMapper_expr = """bill_status_name""".replace("input_file_name()", "input_file") _BillWriterMapper_expr = _BillWriterMapper_expr.replace("_dq_source_file", "input_file") if "." in _BillWriterMapper_expr: _BillWriterMapper_expr = rewrite_response_body_json_access(_BillWriterMapper_expr) _BillWriterMapper_select_clause.append(f"{_BillWriterMapper_expr} AS bill_status_name") _BillWriterMapper_expr = """completion_dttm""".replace("input_file_name()", "input_file") _BillWriterMapper_expr = _BillWriterMapper_expr.replace("_dq_source_file", "input_file") if "." in _BillWriterMapper_expr: _BillWriterMapper_expr = rewrite_response_body_json_access(_BillWriterMapper_expr) _BillWriterMapper_select_clause.append(f"{_BillWriterMapper_expr} AS completion_dttm") _BillWriterMapper_mapping_sql = ("SELECT " + ', '.join(_BillWriterMapper_select_clause) + " FROM filter__1_df").replace("{job_id}", f"'{job_id}'") _BillWriterMapper_input_data = { "component": "BillWriterMapper", "datasource": "filter__1", "include_existing_columns": False, "to_schema_field_count": 11, } try: BillWriterMapper_df = spark.sql(_BillWriterMapper_mapping_sql) except AnalysisException as e: handle_analysis_error( e, component_name="BillWriterMapper", error_code="TRF-MAP-002", exception_class=SchemaMappingException, message=f"Spark analysis error during BillWriterMapper mapping: {e!s}", job_id=job_id, workspace=workspace, workflow=workflow, execution_environment=execution_environment, extra_details={"retry_job_id": retry_job_id or None, "sql_preview": _BillWriterMapper_mapping_sql[:2000]}, input_data=_BillWriterMapper_input_data, ) except Py4JJavaError as e: handle_java_error( e, component_name="BillWriterMapper", operation="mapping SQL", format_name="sql", job_id=job_id, workspace=workspace, workflow=workflow, execution_environment=execution_environment, override_class=ExpressionException, override_code="TRF-EXP-001", extra_details={"retry_job_id": retry_job_id or None, "sql_preview": _BillWriterMapper_mapping_sql[:2000]}, input_data=_BillWriterMapper_input_data, ) BillWriterMapper_df, BillWriterMapper_observer = observe_metrics("BillWriterMapper_df", BillWriterMapper_df) BillWriterMapper_df.createOrReplaceTempView("BillWriterMapper_df") BillWriterMapper_dependency_key="BillWriterMapper" print(filter__1_dependency_key) BillWriterMapper_execute_status="SUCCESS" except Exception as e: BillWriterMapper_error = e log_error(LOGGER, f"Component BillWriterMapper Failed", e, component_name="BillWriterMapper") BillWriterMapper_execute_status="ERROR" raise e finally: BillWriterMapper_end_time=time.time() # %% forecast_insight_code_transform_start_time=time.time() try: import builtins import json as json_lib import traceback from datetime import datetime, timedelta import numpy as np # ===================================================== # CHECK PROPHET # ===================================================== try: from prophet import Prophet PROPHET_AVAILABLE = False # print("Prophet installed") except Exception as e: PROPHET_AVAILABLE = True # print(f"Prophet not available: {e}") # ===================================================== # THRESHOLDS (mirrors BillForecastService class constants) # ===================================================== HIGH_USAGE_THRESHOLD = 15.0 # % increase -> high_usage DROP_THRESHOLD = -15.0 # % decrease -> drop_detected MIN_BILLS_FOR_FILTERING = 3 # minimum bills to apply IQR outlier filtering TREND_DAMPEN = 0.5 # apply only 50% of observed MoM change in fallback # ===================================================== # READ SOURCE # ===================================================== source_df = data_mapper__5_df # print("Input schema:") source_df.printSchema() pdf = ( source_df .select( "account_id", "mapper_bill_id", "bill_date", "amount_value" ) .toPandas() ) pdf["bill_date"] = pd.to_datetime(pdf["bill_date"]) # print("Input rows =", len(pdf)) # ===================================================== # MAPPER FILTER — only keep accounts present in mapper_df # ===================================================== mapper_df = BillWriterMapper_df mapper_pdf = ( mapper_df .select("account_id") .toPandas() ) if mapper_pdf.empty: print("Mapper has no data - returning empty output successfully") pdf = pdf.iloc[0:0] else: mapper_account_ids = set(mapper_pdf["account_id"].dropna().unique()) print("Mapper account count =", len(mapper_account_ids)) before_count = len(pdf) pdf = pdf[pdf["account_id"].isin(mapper_account_ids)].reset_index(drop=True) print(f"Filtered source rows by mapper: {before_count} -> {len(pdf)}") print("Input rows after mapper filter =", len(pdf)) output_rows = [] # ===================================================== # HELPER FUNCTIONS — outlier filtering & weighting # ===================================================== def filter_outliers(amounts): """Remove outliers using IQR method. Returns filtered list (at least 2 values kept).""" if len(amounts) < 3: return amounts sorted_vals = sorted(amounts) n = len(sorted_vals) q1 = sorted_vals[n // 4] q3 = sorted_vals[(3 * n) // 4] iqr = q3 - q1 # Use 1.5x IQR rule; if IQR is 0, fall back to median +/- band if iqr > 0: lower_bound = q1 - 1.5 * iqr upper_bound = q3 + 1.5 * iqr else: median = sorted_vals[n // 2] lower_bound = median * 0.2 upper_bound = median * 3.0 filtered = [a for a in amounts if lower_bound <= a <= upper_bound] # Always keep at least the 2 most recent values if len(filtered) < 2: filtered = amounts[-2:] return filtered def exponential_weights(n, decay=0.5): """Generate exponential decay weights - most recent gets highest weight. Example with n=3, decay=0.5: [0.25, 0.5, 1.0] -> normalized to [0.143, 0.286, 0.571] """ raw = [decay ** (n - 1 - i) for i in range(n)] total = builtins.sum(raw) return [w / total for w in raw] def detect_anomalies(amounts, threshold=2.0): """Detect anomalies using Z-score method.""" if len(amounts) < 3: return [False] * len(amounts) mean_val = np.mean(amounts) std_val = np.std(amounts) if std_val == 0: return [False] * len(amounts) z_scores = [(x - mean_val) / std_val for x in amounts] return [bool(abs(z) > threshold) for z in z_scores] def calculate_trend_slope(amounts): """Calculate normalized trend slope using linear regression (% change per period).""" if len(amounts) < 2: return 0.0 x = np.arange(len(amounts)) y = np.array(amounts) n = len(x) denom = (n * np.sum(x ** 2) - np.sum(x) ** 2) if denom == 0: return 0.0 slope = (n * np.sum(x * y) - np.sum(x) * np.sum(y)) / denom mean_val = np.mean(amounts) if mean_val > 0: return (slope / mean_val) * 100 return 0.0 # ===================================================== # HELPER FUNCTIONS — forecasting # ===================================================== def prophet_forecast(prophet_df, periods=3): """Run Prophet forecast for the next `periods` months. Raises on failure so the caller can fall back to fallback_forecast_weighted.""" model = Prophet( yearly_seasonality=True, weekly_seasonality=False, daily_seasonality=False, interval_width=0.80 ) model.fit(prophet_df) future = model.make_future_dataframe( periods=periods, freq="M" ) pred = model.predict(future) forecast_df = pred[ pred["ds"] > prophet_df["ds"].max() ][[ "ds", "yhat", "yhat_lower", "yhat_upper" ]].copy() forecast_df["yhat"] = forecast_df["yhat"].clip(lower=0).round(2) forecast_df["yhat_lower"] = forecast_df["yhat_lower"].clip(lower=0).round(2) forecast_df["yhat_upper"] = forecast_df["yhat_upper"].round(2) return forecast_df def fallback_forecast_weighted(prophet_df, periods=3): """ Weighted-average fallback with dampened trend (used when Prophet is unavailable, fails, or there are fewer than 4 data points). Steps: 1. Outlier filtering (IQR method) when >= MIN_BILLS_FOR_FILTERING bills. 2. Exponential decay weighting (most recent bill weighted highest). 3. Dampened month-over-month trend projection (50% of observed rate), with a seasonal override when a same-calendar-month average exists. Confidence interval: +/-15% around the forecast value. """ amounts_all = [float(v) for v in prophet_df["y"].values] dates_all = list(prophet_df["ds"].values) amounts = [a for a in amounts_all if a > 0] if not amounts: return pd.DataFrame(columns=["ds", "yhat", "yhat_lower", "yhat_upper"]) # Step 1: outlier filtering if len(amounts) >= MIN_BILLS_FOR_FILTERING: clean_amounts = filter_outliers(amounts) else: clean_amounts = amounts # Seasonal map: month-of-year -> list of historical amounts in that month monthly_map = {} for d, a in zip(dates_all, amounts_all): if a <= 0: continue month = pd.Timestamp(d).month monthly_map.setdefault(month, []).append(a) last_date = prophet_df["ds"].max() # Step 2: exponential decay weighted average on clean data weights = exponential_weights(len(clean_amounts)) weighted_avg = builtins.sum(a * w for a, w in zip(clean_amounts, weights)) # Step 3: dampened month-over-month trend if len(clean_amounts) >= 2: mom_changes = [] for j in range(1, len(clean_amounts)): if clean_amounts[j - 1] > 0: mom_changes.append( (clean_amounts[j] - clean_amounts[j - 1]) / clean_amounts[j - 1] ) avg_mom = (builtins.sum(mom_changes) / len(mom_changes)) if mom_changes else 0.0 dampened_mom = avg_mom * TREND_DAMPEN else: dampened_mom = 0.0 rows = [] base_val = weighted_avg for i in range(1, periods + 1): future_dt = last_date + pd.DateOffset(months=i) future_month = future_dt.month if future_month in monthly_map and monthly_map[future_month]: seasonal_avg = builtins.sum(monthly_map[future_month]) / len(monthly_map[future_month]) predicted_value = seasonal_avg else: predicted_value = builtins.max(0.0, base_val * (1 + dampened_mom) ** i) lower_bound = builtins.max(0.0, predicted_value * 0.85) upper_bound = predicted_value * 1.15 rows.append({ "ds": future_dt, "yhat": round(predicted_value, 2), "yhat_lower": round(lower_bound, 2), "yhat_upper": round(upper_bound, 2) }) # print( # f"Fallback forecast: {len(amounts)} bills -> {len(clean_amounts)} clean -> " # f"base ${weighted_avg:.2f}, dampened MoM {dampened_mom * 100:.1f}%" # ) return pd.DataFrame(rows) # ===================================================== # HELPER FUNCTIONS — classification, severity, explanation # ===================================================== def classify_type(recent_amounts, forecast_amounts): """Classify insight type based on % change between recent avg and forecast avg.""" if not recent_amounts or not forecast_amounts: return "stable_usage" recent_avg = builtins.sum(recent_amounts) / len(recent_amounts) forecast_avg = builtins.sum(forecast_amounts) / len(forecast_amounts) if recent_avg == 0: return "stable_usage" pct_change = ((forecast_avg - recent_avg) / recent_avg) * 100 if pct_change >= HIGH_USAGE_THRESHOLD: return "high_usage" elif pct_change <= DROP_THRESHOLD: return "drop_detected" else: return "stable_usage" def compute_severity_score(recent_amounts, forecast_amounts): """Compute a 1-10 severity score based on magnitude of change.""" if not recent_amounts or not forecast_amounts: return 1 recent_avg = builtins.sum(recent_amounts) / len(recent_amounts) forecast_avg = builtins.sum(forecast_amounts) / len(forecast_amounts) if recent_avg == 0: return 1 pct_change = abs(((forecast_avg - recent_avg) / recent_avg) * 100) return builtins.min(10, builtins.max(1, int(pct_change / 10) + 1)) def generate_explanation_template(recent_amounts, forecast_amounts, insight_type, bill_count): """Template-based alert/message/explanation for the Rank 1 forecast insight.""" recent_avg = builtins.sum(recent_amounts) / len(recent_amounts) if recent_amounts else 0 forecast_avg = builtins.sum(forecast_amounts) / len(forecast_amounts) if forecast_amounts else 0 if recent_avg > 0: pct_change = ((forecast_avg - recent_avg) / recent_avg) * 100 else: pct_change = 0 direction = "increase" if pct_change > 0 else "decrease" abs_pct = abs(pct_change) type_labels = { "high_usage": "High Usage Expected", "drop_detected": "Bill Drop Detected", "stable_usage": "Stable Billing Pattern", } alert = type_labels.get(insight_type, "Bill Forecast") message = f"Forecasted bills show a {abs_pct:.1f}% {direction} over the next 3 months." explanation = ( f"Based on the last {bill_count} months of billing data, " f"the average recent bill is ${recent_avg:.2f} and the forecasted average is ${forecast_avg:.2f}. " f"This represents a {abs_pct:.1f}% {direction} " f"(${abs(forecast_avg - recent_avg):.2f} difference)." ) return alert, message, explanation # ===================================================== # HELPER FUNCTIONS — graphs & considered bills # ===================================================== def build_bar_graph(labels, data, label, color="rgba(75,192,192,0.6)"): return { "type": "bar", "labels": labels, "datasets": [{ "label": label, "data": data, "backgroundColor": color }] } def build_forecast_graph(forecast_labels, forecast_values, forecast_lower, forecast_upper): return { "type": "bar", "labels": forecast_labels, "datasets": [ { "label": "Forecasted Bills", "data": forecast_values, "borderColor": "rgb(75,102,192)", "backgroundColor": "rgba(75,192,192,0.2)", "fill": True }, { "label": "Confidence Lower", "data": forecast_lower, "borderColor": "rgba(75,192,192,0.3)", "backgroundColor": "transparent", "borderDash": [5, 5], "fill": False }, { "label": "Confidence Upper", "data": forecast_upper, "borderColor": "rgba(75,192,192,0.3)", "backgroundColor": "transparent", "borderDash": [5, 5], "fill": False } ] } def build_considered_bills(rows_df, anomaly_flags=None): """Build the consideredBills list from a pandas slice of bill rows.""" considered = [] for i, (_, r) in enumerate(rows_df.iterrows()): is_anomaly = bool(anomaly_flags[i]) if anomaly_flags and i < len(anomaly_flags) else False considered.append({ "billID": str(r["mapper_bill_id"]), "billDate": r["bill_date"].strftime("%Y-%m-%d"), "billAmount": round(float(r["amount_value"]), 2), "consumptionValue": round(float(r["amount_value"]), 2), "consumptionUnit": "USD", "isAnomaly": is_anomaly }) return considered # ===================================================== # INSIGHT BUILDER — Rank 1: Bill Forecast # ===================================================== def build_forecast_insight(recent, forecast_df, anomaly_flags): """ Rank 1 insight: forecast classification (high_usage / drop_detected / stable_usage), with historical graph + forecast graph + consideredBills. """ actual_labels = [d.strftime("%Y-%m") for d in recent["bill_date"]] actual_amounts = [round(float(x), 2) for x in recent["amount_value"]] forecast_labels = [d.strftime("%Y-%m") for d in forecast_df["ds"]] forecast_values = [round(float(x), 2) for x in forecast_df["yhat"]] forecast_lower = [round(float(x), 2) for x in forecast_df["yhat_lower"]] forecast_upper = [round(float(x), 2) for x in forecast_df["yhat_upper"]] insight_type = classify_type(actual_amounts, forecast_values) severity = compute_severity_score(actual_amounts, forecast_values) alert, message, explanation = generate_explanation_template( actual_amounts, forecast_values, insight_type, len(recent) ) considered_bills = build_considered_bills(recent, anomaly_flags) actual_graph = build_bar_graph(actual_labels, actual_amounts, "Bills USD") forecast_graph = build_forecast_graph(forecast_labels, forecast_values, forecast_lower, forecast_upper) insight = { "rank": 1, "alert": alert, "message": message, "explanation": explanation, "severityScore": severity, "consideredBills": considered_bills, "graph": actual_graph, "forecastGraph": forecast_graph, "type": insight_type } return insight, forecast_labels, forecast_values, forecast_lower, forecast_upper # ===================================================== # INSIGHT BUILDER — Rank 2: Trend Summary (last 3 months) # ===================================================== def build_trend_insight(recent, anomaly_flags, forecast_labels, forecast_values, forecast_lower, forecast_upper): """ Rank 2 insight: month-over-month trend pattern across the last 3 bills. Patterns: declining_trend - both MoM changes < -20% increasing_trend - both MoM changes > +20% spike_resolved - oldest month 30%+ higher, bills dropped since mid_spike - middle month 30%+ higher than neighbors recent_spike - most recent month jumped 30%+ stable_trend - all within 20% of 3-month average Returns None if fewer than 3 bills or no clear pattern. """ if len(recent) < 3: return None amounts = [round(float(x), 2) for x in recent["amount_value"]] dates = [d.strftime("%Y-%m") for d in recent["bill_date"]] month_names = [d.strftime("%B %Y") for d in recent["bill_date"]] a0, a1, a2 = amounts # oldest -> newest def pct(old, new): return ((new - old) / old * 100) if old != 0 else 0 chg_1 = pct(a0, a1) chg_2 = pct(a1, a2) total_chg = pct(a0, a2) peak_idx = amounts.index(builtins.max(amounts)) alert = "" message = "" explanation = "" trend_type = "stable_trend" if chg_1 < -20 and chg_2 < -20: trend_type = "declining_trend" alert = "Bills Declining Steadily" message = ( f"Your bill has dropped {abs(total_chg):.0f}% over the last 3 months " f"- from ${a0:,.2f} in {month_names[0]} to ${a2:,.2f} in {month_names[2]}." ) explanation = ( f"{month_names[0]}: ${a0:,.2f} -> {month_names[1]}: ${a1:,.2f} ({chg_1:+.1f}%) -> " f"{month_names[2]}: ${a2:,.2f} ({chg_2:+.1f}%). " f"This is a consistent downward trend that may continue." ) elif chg_1 > 20 and chg_2 > 20: trend_type = "increasing_trend" alert = "Bills Increasing Steadily" message = ( f"Your bill has risen {abs(total_chg):.0f}% over the last 3 months " f"- from ${a0:,.2f} in {month_names[0]} to ${a2:,.2f} in {month_names[2]}." ) explanation = ( f"{month_names[0]}: ${a0:,.2f} -> {month_names[1]}: ${a1:,.2f} ({chg_1:+.1f}%) -> " f"{month_names[2]}: ${a2:,.2f} ({chg_2:+.1f}%). " f"Your usage has been climbing; consider reviewing recent activity." ) elif peak_idx == 0 and abs(total_chg) > 30: trend_type = "spike_resolved" alert = "Recent Bill Spike Has Resolved" message = ( f"Your bill was ${a0:,.2f} in {month_names[0]} but has since dropped " f"to ${a2:,.2f} in {month_names[2]}." ) explanation = ( f"The {month_names[0]} bill (${a0:,.2f}) was significantly higher than the recent " f"{month_names[1]} (${a1:,.2f}) and {month_names[2]} (${a2:,.2f}). " f"This suggests the spike was a one-time event and bills are normalizing." ) elif peak_idx == 1 and pct(a1, a0) < -30 and pct(a1, a2) < -30: trend_type = "mid_spike" alert = f"Bill Spike in {month_names[1]}" message = ( f"Your {month_names[1]} bill spiked to ${a1:,.2f} but has returned " f"to ${a2:,.2f} in {month_names[2]}." ) explanation = ( f"{month_names[0]}: ${a0:,.2f} -> {month_names[1]}: ${a1:,.2f} " f"(spike of {pct(a0, a1):+.1f}%) -> " f"{month_names[2]}: ${a2:,.2f} (back to {pct(a1, a2):+.1f}%). " f"The {month_names[1]} spike appears to be an anomaly." ) elif peak_idx == 2 and pct(a1, a2) > 30: trend_type = "recent_spike" alert = "Recent Bill Spike" message = ( f"Your latest bill in {month_names[2]} jumped to ${a2:,.2f} " f"- up {pct(a1, a2):.0f}% from {month_names[1]} (${a1:,.2f})." ) explanation = ( f"{month_names[0]}: ${a0:,.2f} -> {month_names[1]}: ${a1:,.2f} -> " f"{month_names[2]}: ${a2:,.2f} ({pct(a1, a2):+.1f}%). " f"This recent increase is worth monitoring." ) else: avg_3 = builtins.sum(amounts) / 3 max_dev = builtins.max(abs(a - avg_3) / avg_3 * 100 for a in amounts) if avg_3 > 0 else 0 if max_dev < 20: trend_type = "stable_trend" alert = "Bills Are Stable" message = f"Your bills have been consistent over the last 3 months, averaging ${avg_3:,.2f}." explanation = ( f"{month_names[0]}: ${a0:,.2f}, {month_names[1]}: ${a1:,.2f}, {month_names[2]}: ${a2:,.2f}. " f"Variation is within normal range." ) else: return None severity = builtins.min(10, builtins.max(1, int(abs(total_chg) / 15) + 1)) considered_bills = build_considered_bills(recent, anomaly_flags) trend_graph = { "type": "bar", "labels": dates, "datasets": [{ "label": "Monthly Bills USD", "data": amounts, "backgroundColor": [ "rgba(255,99,132,0.6)" if i == peak_idx else "rgba(75,192,192,0.6)" for i in range(3) ] }] } forecast_graph = build_forecast_graph(forecast_labels, forecast_values, forecast_lower, forecast_upper) return { "rank": 2, "alert": alert, "message": message, "explanation": explanation, "severityScore": severity, "consideredBills": considered_bills, "graph": trend_graph, "forecastGraph": forecast_graph, "type": trend_type } # ===================================================== # INSIGHT BUILDER — Rank 3: Year-over-Year Comparison # ===================================================== def build_yoy_insight(acct_df, anomaly_flags_full): """ Rank 3 insight: compares the most recent 3 months against the same 3 calendar months a year ago. Requires >= 6 months of history overall, and requires that the same-month-last-year data actually exists. Returns None if insufficient data. """ if len(acct_df) < 6: return None sorted_df = acct_df.sort_values("bill_date").reset_index(drop=True) recent_3 = sorted_df.tail(3) recent_dates = [d for d in recent_3["bill_date"]] yoy_targets = [(d.year - 1, d.month) for d in recent_dates] yoy_rows = sorted_df[ sorted_df["bill_date"].apply(lambda d: (d.year, d.month) in yoy_targets) ] if len(yoy_rows) < len(yoy_targets): return None yoy_rows = yoy_rows.sort_values("bill_date").tail(len(yoy_targets)) recent_amounts = [round(float(x), 2) for x in recent_3["amount_value"]] yoy_amounts = [round(float(x), 2) for x in yoy_rows["amount_value"]] recent_avg = builtins.sum(recent_amounts) / len(recent_amounts) yoy_avg = builtins.sum(yoy_amounts) / len(yoy_amounts) pct_change = ((recent_avg - yoy_avg) / yoy_avg * 100) if yoy_avg != 0 else 0 direction = "increased" if pct_change > 0 else "decreased" if pct_change > HIGH_USAGE_THRESHOLD: insight_type = "high_usage" elif pct_change < DROP_THRESHOLD: insight_type = "drop_detected" else: insight_type = "stable_usage" severity = builtins.min(10, builtins.max(1, int(abs(pct_change) / 10) + 1)) # anomaly flags computed over the full account history align by position recent_anomalies = anomaly_flags_full[-3:] if len(anomaly_flags_full) >= 3 else [False] * 3 considered_bills = build_considered_bills(recent_3, recent_anomalies) + build_considered_bills(yoy_rows, None) yoy_graph = build_bar_graph( [d.strftime("%Y-%m") for d in yoy_rows["bill_date"]], yoy_amounts, "Last Year Bills USD", color="rgba(153,102,255,0.6)" ) current_graph = build_bar_graph( [d.strftime("%Y-%m") for d in recent_3["bill_date"]], recent_amounts, "Current Year Bills USD", color="rgba(75,192,192,0.6)" ) return { "rank": 3, "alert": f"Year-over-Year Bill {direction.capitalize()}", "message": f"Your bills have {direction} by {abs(pct_change):.1f}% compared to the same period last year.", "explanation": ( f"Average bill for the recent 3 months: ${recent_avg:.2f}. " f"Average bill for the same 3 months last year: ${yoy_avg:.2f}. " f"That is a {abs(pct_change):.1f}% {direction}." ), "severityScore": severity, "consideredBills": considered_bills, "graph": yoy_graph, "forecastGraph": current_graph, "type": insight_type } # ===================================================== # FORECAST ACCURACY — best-effort in-sample backtest # ===================================================== def compute_forecast_accuracy_backtest(acct_df): """ Best-effort forecast accuracy, computed entirely from the bills already present in this dataframe (no external previous-forecast input available). Approach: hold out the most recent actual bill, forecast 1 month ahead using only the months before it (same Prophet/fallback logic as the live forecast), then compare that 1-month-ahead prediction against the real bill that came in. This approximates "how accurate was last month's forecast" without needing a stored previous forecast. Requires >= 4 bills (3 to forecast from + 1 actual to validate against). Returns None if not enough data. """ sorted_df = acct_df.sort_values("bill_date").reset_index(drop=True) if len(sorted_df) < 4: return None train_df = sorted_df.iloc[:-1] actual_row = sorted_df.iloc[-1] train_prophet_df = train_df.rename(columns={"bill_date": "ds", "amount_value": "y"})[["ds", "y"]] try: if PROPHET_AVAILABLE and len(train_df) >= 4: bt_forecast_df = prophet_forecast(train_prophet_df, periods=1) else: raise Exception("Prophet unavailable or insufficient data for backtest") except Exception: bt_forecast_df = fallback_forecast_weighted(train_prophet_df, periods=1) if bt_forecast_df.empty: return None predicted = float(bt_forecast_df.iloc[0]["yhat"]) actual = float(actual_row["amount_value"]) error_pct = abs(predicted - actual) / actual * 100 if actual != 0 else 0 accuracy = builtins.max(0, 100 - error_pct) return { "method": "in_sample_backtest", "validatedMonth": actual_row["bill_date"].strftime("%Y-%m"), "predicted": round(predicted, 2), "actual": round(actual, 2), "accuracyPct": round(accuracy, 1) } # ===================================================== # PROCESS EACH ACCOUNT # ===================================================== for account_id, acct_df in pdf.groupby("account_id"): try: # print(f"\nProcessing account {account_id}") acct_df = acct_df.sort_values("bill_date").reset_index(drop=True) bill_count = len(acct_df) # print("Bill count =", bill_count) if bill_count < 3: # print("Skipping account - less than 3 bills") continue # ============================================= # PROPHET INPUT # ============================================= prophet_df = acct_df.rename( columns={ "bill_date": "ds", "amount_value": "y" } )[["ds", "y"]] # ============================================= # FORECAST (Prophet >= 4 points, else weighted fallback) # ============================================= try: if PROPHET_AVAILABLE and bill_count >= 4: # print("Running Prophet") forecast_df = prophet_forecast(prophet_df, periods=3) else: raise Exception("Prophet unavailable or insufficient data (<4 points)") except Exception as prophet_error: # print(f"Prophet failed for {account_id}: {prophet_error}") # print("Using weighted-average fallback (outlier filtering + exponential decay + dampened trend)") forecast_df = fallback_forecast_weighted(prophet_df, periods=3) # print("Forecast rows =", len(forecast_df)) # print("forecast_df", forecast_df) # ============================================= # ANOMALY DETECTION (full history, Z-score) # ============================================= all_amounts = [round(float(x), 2) for x in acct_df["amount_value"]] anomaly_flags_full = detect_anomalies(all_amounts) # ============================================= # ACTUAL DATA — recent 3 months # ============================================= recent = acct_df.tail(3).reset_index(drop=True) recent_anomaly_flags = anomaly_flags_full[-3:] if len(anomaly_flags_full) >= 3 else [False] * len(recent) # print("recent", recent) # ============================================= # TREND SLOPE (informational, kept in output) # ============================================= recent_amounts_for_slope = [round(float(x), 2) for x in recent["amount_value"]] trend_slope = calculate_trend_slope(recent_amounts_for_slope) # ============================================= # RANK 1 — FORECAST INSIGHT # ============================================= forecast_insight, forecast_labels, forecast_values, forecast_lower, forecast_upper = ( build_forecast_insight(recent, forecast_df, recent_anomaly_flags) ) insights = [forecast_insight] # ============================================= # RANK 2 — TREND SUMMARY # ============================================= trend_insight = build_trend_insight( recent, recent_anomaly_flags, forecast_labels, forecast_values, forecast_lower, forecast_upper ) if trend_insight: insights.append(trend_insight) # ============================================= # RANK 3 — YEAR-OVER-YEAR COMPARISON # ============================================= yoy_insight = build_yoy_insight(acct_df, anomaly_flags_full) if yoy_insight: insights.append(yoy_insight) # ============================================= # FORECAST ACCURACY — best-effort backtest # ============================================= accuracy_data = compute_forecast_accuracy_backtest(acct_df) if accuracy_data: for ins in insights: if ins.get("rank") == 1: ins["previousForecastAccuracy"] = accuracy_data # ============================================= # FINAL JSON # ============================================= response = { "accountId": account_id, "generatedAt": datetime.utcnow().isoformat(), "cacheHit": False, "dataPointsUsed": len(recent), "nextRefreshDate": ( datetime.utcnow() + pd.DateOffset(months=1) ).strftime("%Y-%m-%d"), "forecastMonth": datetime.utcnow().strftime("%Y-%m"), "stale": False, "forecastMethod": "prophet" if (PROPHET_AVAILABLE and bill_count >= 4) else "weighted_average_fallback", "trendSlope": round(trend_slope, 2), "insights": insights } output_rows.append( Row( account_id=str(account_id), insight_json=json_lib.dumps(response) ) ) except Exception as account_error: # print(f"Account failed {account_id}: {account_error}") traceback.print_exc() continue # ===================================================== # OUTPUT # ===================================================== if len(output_rows) > 0: forecast_insight_code_transform_df = ( spark.createDataFrame(output_rows) ) else: empty_schema = ( "account_id string," " insight_json string" ) forecast_insight_code_transform_df = ( spark.createDataFrame( [], empty_schema ) ) # print(f"Generated insights for {len(output_rows)} accounts") # forecast_insight_code_transform_df.show(truncate=False) forecast_insight_code_transform_df.createOrReplaceTempView("forecast_insight_code_transform_df") forecast_insight_code_transform_execute_status = "SUCCESS" except Exception as e: print("Pipeline failed") print(str(e)) forecast_insight_code_transform_execute_status = "ERROR" raise forecast_insight_code_transform_end_time=time.time() forecast_insight_code_transform_dependency_key="forecast_insight_code_transform" print(BillWriterMapper_dependency_key) print(data_mapper__5_dependency_key) # %% forecast_insight_data_mapper_start_time=time.time() forecast_insight_data_mapper_fail_on_error="True" try: _forecast_insight_data_mapper_select_clause=[] _forecast_insight_data_mapper_expr = """account_id""".replace("input_file_name()", "input_file") _forecast_insight_data_mapper_expr = _forecast_insight_data_mapper_expr.replace("_dq_source_file", "input_file") if "." in _forecast_insight_data_mapper_expr: _forecast_insight_data_mapper_expr = rewrite_response_body_json_access(_forecast_insight_data_mapper_expr) _forecast_insight_data_mapper_select_clause.append(f"{_forecast_insight_data_mapper_expr} AS account_id") _forecast_insight_data_mapper_expr = """uuid()""".replace("input_file_name()", "input_file") _forecast_insight_data_mapper_expr = _forecast_insight_data_mapper_expr.replace("_dq_source_file", "input_file") if "." in _forecast_insight_data_mapper_expr: _forecast_insight_data_mapper_expr = rewrite_response_body_json_access(_forecast_insight_data_mapper_expr) _forecast_insight_data_mapper_select_clause.append(f"{_forecast_insight_data_mapper_expr} AS id") _forecast_insight_data_mapper_expr = """get_json_object(insight_json, \'$.insights\')""".replace("input_file_name()", "input_file") _forecast_insight_data_mapper_expr = _forecast_insight_data_mapper_expr.replace("_dq_source_file", "input_file") if "." in _forecast_insight_data_mapper_expr: _forecast_insight_data_mapper_expr = rewrite_response_body_json_access(_forecast_insight_data_mapper_expr) _forecast_insight_data_mapper_select_clause.append(f"{_forecast_insight_data_mapper_expr} AS insight_data") _forecast_insight_data_mapper_expr = """get_json_object(insight_json, \'$.forecastMonth\')""".replace("input_file_name()", "input_file") _forecast_insight_data_mapper_expr = _forecast_insight_data_mapper_expr.replace("_dq_source_file", "input_file") if "." in _forecast_insight_data_mapper_expr: _forecast_insight_data_mapper_expr = rewrite_response_body_json_access(_forecast_insight_data_mapper_expr) _forecast_insight_data_mapper_select_clause.append(f"{_forecast_insight_data_mapper_expr} AS forecast_month") _forecast_insight_data_mapper_expr = """get_json_object(insight_json, \'$.insights[0].type\')""".replace("input_file_name()", "input_file") _forecast_insight_data_mapper_expr = _forecast_insight_data_mapper_expr.replace("_dq_source_file", "input_file") if "." in _forecast_insight_data_mapper_expr: _forecast_insight_data_mapper_expr = rewrite_response_body_json_access(_forecast_insight_data_mapper_expr) _forecast_insight_data_mapper_select_clause.append(f"{_forecast_insight_data_mapper_expr} AS type") _forecast_insight_data_mapper_expr = """date_format(current_timestamp(), "yyyy-MM-dd\'T\'HH:mm:ss")""".replace("input_file_name()", "input_file") _forecast_insight_data_mapper_expr = _forecast_insight_data_mapper_expr.replace("_dq_source_file", "input_file") if "." in _forecast_insight_data_mapper_expr: _forecast_insight_data_mapper_expr = rewrite_response_body_json_access(_forecast_insight_data_mapper_expr) _forecast_insight_data_mapper_select_clause.append(f"{_forecast_insight_data_mapper_expr} AS created_at") _forecast_insight_data_mapper_expr = """date_format(current_timestamp(), "yyyy-MM-dd\'T\'HH:mm:ss")""".replace("input_file_name()", "input_file") _forecast_insight_data_mapper_expr = _forecast_insight_data_mapper_expr.replace("_dq_source_file", "input_file") if "." in _forecast_insight_data_mapper_expr: _forecast_insight_data_mapper_expr = rewrite_response_body_json_access(_forecast_insight_data_mapper_expr) _forecast_insight_data_mapper_select_clause.append(f"{_forecast_insight_data_mapper_expr} AS updated_at") _forecast_insight_data_mapper_mapping_sql = ("SELECT " + ', '.join(_forecast_insight_data_mapper_select_clause) + " FROM forecast_insight_code_transform_df").replace("{job_id}", f"'{job_id}'") _forecast_insight_data_mapper_input_data = { "component": "forecast_insight_data_mapper", "datasource": "forecast_insight_code_transform", "include_existing_columns": False, "to_schema_field_count": 7, } try: forecast_insight_data_mapper_df = spark.sql(_forecast_insight_data_mapper_mapping_sql) except AnalysisException as e: handle_analysis_error( e, component_name="forecast_insight_data_mapper", error_code="TRF-MAP-002", exception_class=SchemaMappingException, message=f"Spark analysis error during forecast_insight_data_mapper mapping: {e!s}", job_id=job_id, workspace=workspace, workflow=workflow, execution_environment=execution_environment, extra_details={"retry_job_id": retry_job_id or None, "sql_preview": _forecast_insight_data_mapper_mapping_sql[:2000]}, input_data=_forecast_insight_data_mapper_input_data, ) except Py4JJavaError as e: handle_java_error( e, component_name="forecast_insight_data_mapper", operation="mapping SQL", format_name="sql", job_id=job_id, workspace=workspace, workflow=workflow, execution_environment=execution_environment, override_class=ExpressionException, override_code="TRF-EXP-001", extra_details={"retry_job_id": retry_job_id or None, "sql_preview": _forecast_insight_data_mapper_mapping_sql[:2000]}, input_data=_forecast_insight_data_mapper_input_data, ) forecast_insight_data_mapper_df, forecast_insight_data_mapper_observer = observe_metrics("forecast_insight_data_mapper_df", forecast_insight_data_mapper_df) forecast_insight_data_mapper_df.createOrReplaceTempView("forecast_insight_data_mapper_df") forecast_insight_data_mapper_dependency_key="forecast_insight_data_mapper" print(forecast_insight_code_transform_dependency_key) forecast_insight_data_mapper_execute_status="SUCCESS" except Exception as e: forecast_insight_data_mapper_error = e log_error(LOGGER, f"Component forecast_insight_data_mapper Failed", e, component_name="forecast_insight_data_mapper") forecast_insight_data_mapper_execute_status="ERROR" raise e finally: forecast_insight_data_mapper_end_time=time.time() # %% data_writer__1_start_time=time.time() data_writer__1_fail_on_error="" try: _data_writer__1_options = { 'jdbc':{ 'dbtable': 'bills', 'url':secrets.get(''), 'driver':'', 'stringtype': 'unspecified' }, 'kafka' : { 'kafka.bootstrap.servers' : secrets.get('OCULAR_KAFKA_BOOTSTRAP_SERVERS'), 'topic' : '' } } _BillWriterMapper_df = BillWriterMapper_df _data_writer__1_writer = _BillWriterMapper_df.write.format('iceberg').mode('append') _data_writer__1_save_path = 'dremio.bills' _data_writer__1_write_input = { "component": "data_writer__1", "datasource": "BillWriterMapper", "iceberg_catalog": "dremio", "table_name": "bills", } try: _data_writer__1_writer.save(_data_writer__1_save_path) except AnalysisException as e: handle_analysis_error( e, component_name="data_writer__1", message=f"Write failed for 'data_writer__1' to {_data_writer__1_save_path or 'iceberg'}: {e!s}", job_id=job_id, workspace=workspace, workflow=workflow, execution_environment=execution_environment, extra_details={"format": "iceberg", "mode": "append", "save_path": _data_writer__1_save_path}, input_data=_data_writer__1_write_input, ) except Py4JJavaError as e: handle_java_error( e, component_name="data_writer__1", operation="write", format_name="iceberg", path=_data_writer__1_save_path, job_id=job_id, workspace=workspace, workflow=workflow, execution_environment=execution_environment, extra_details={"mode": "append"}, input_data=_data_writer__1_write_input, ) data_writer__1_dependency_key="data_writer__1" print(BillWriterMapper_dependency_key) print(forecast_insight_data_mapper_dependency_key) data_writer__1_execute_status="SUCCESS" except Exception as e: data_writer__1_error = e log_error(LOGGER, f"Component data_writer__1 Failed", e, component_name="data_writer__1") data_writer__1_execute_status="ERROR" raise e finally: data_writer__1_end_time=time.time() # %% accountinsights_data_writer_start_time=time.time() accountinsights_data_writer_fail_on_error="" try: _accountinsights_data_writer_options = { 'jdbc':{ 'dbtable': 'insights', 'url':secrets.get(''), 'driver':'', 'stringtype': 'unspecified' }, 'kafka' : { 'kafka.bootstrap.servers' : secrets.get('OCULAR_KAFKA_BOOTSTRAP_SERVERS'), 'topic' : '' } } _forecast_insight_data_mapper_df = forecast_insight_data_mapper_df _accountinsights_data_writer_writer = _forecast_insight_data_mapper_df.write.format('iceberg').mode('append') _accountinsights_data_writer_save_path = 'dremio.insights' _accountinsights_data_writer_write_input = { "component": "accountinsights_data_writer", "datasource": "forecast_insight_data_mapper", "iceberg_catalog": "dremio", "table_name": "insights", } try: _accountinsights_data_writer_writer.save(_accountinsights_data_writer_save_path) except AnalysisException as e: handle_analysis_error( e, component_name="accountinsights_data_writer", message=f"Write failed for 'accountinsights_data_writer' to {_accountinsights_data_writer_save_path or 'iceberg'}: {e!s}", job_id=job_id, workspace=workspace, workflow=workflow, execution_environment=execution_environment, extra_details={"format": "iceberg", "mode": "append", "save_path": _accountinsights_data_writer_save_path}, input_data=_accountinsights_data_writer_write_input, ) except Py4JJavaError as e: handle_java_error( e, component_name="accountinsights_data_writer", operation="write", format_name="iceberg", path=_accountinsights_data_writer_save_path, job_id=job_id, workspace=workspace, workflow=workflow, execution_environment=execution_environment, extra_details={"mode": "append"}, input_data=_accountinsights_data_writer_write_input, ) accountinsights_data_writer_dependency_key="accountinsights_data_writer" print(forecast_insight_data_mapper_dependency_key) accountinsights_data_writer_execute_status="SUCCESS" except Exception as e: accountinsights_data_writer_error = e log_error(LOGGER, f"Component accountinsights_data_writer Failed", e, component_name="accountinsights_data_writer") accountinsights_data_writer_execute_status="ERROR" raise e finally: accountinsights_data_writer_end_time=time.time() # %% customerMapper_start_time=time.time() customerMapper_fail_on_error="True" try: _customerMapper_select_clause=[] _customerMapper_expr = """account_id""".replace("input_file_name()", "input_file") _customerMapper_expr = _customerMapper_expr.replace("_dq_source_file", "input_file") if "." in _customerMapper_expr: _customerMapper_expr = rewrite_response_body_json_access(_customerMapper_expr) _customerMapper_select_clause.append(f"{_customerMapper_expr} AS account_id") _customerMapper_expr = """current_timestamp()""".replace("input_file_name()", "input_file") _customerMapper_expr = _customerMapper_expr.replace("_dq_source_file", "input_file") if "." in _customerMapper_expr: _customerMapper_expr = rewrite_response_body_json_access(_customerMapper_expr) _customerMapper_select_clause.append(f"{_customerMapper_expr} AS created_at") _customerMapper_expr = """id""".replace("input_file_name()", "input_file") _customerMapper_expr = _customerMapper_expr.replace("_dq_source_file", "input_file") if "." in _customerMapper_expr: _customerMapper_expr = rewrite_response_body_json_access(_customerMapper_expr) _customerMapper_select_clause.append(f"{_customerMapper_expr} AS id") _customerMapper_expr = """bill_status""".replace("input_file_name()", "input_file") _customerMapper_expr = _customerMapper_expr.replace("_dq_source_file", "input_file") if "." in _customerMapper_expr: _customerMapper_expr = rewrite_response_body_json_access(_customerMapper_expr) _customerMapper_select_clause.append(f"{_customerMapper_expr} AS status") _customerMapper_expr = """bill_date""".replace("input_file_name()", "input_file") _customerMapper_expr = _customerMapper_expr.replace("_dq_source_file", "input_file") if "." in _customerMapper_expr: _customerMapper_expr = rewrite_response_body_json_access(_customerMapper_expr) _customerMapper_select_clause.append(f"{_customerMapper_expr} AS latest_bill_date") _customerMapper_expr = """bill_id""".replace("input_file_name()", "input_file") _customerMapper_expr = _customerMapper_expr.replace("_dq_source_file", "input_file") if "." in _customerMapper_expr: _customerMapper_expr = rewrite_response_body_json_access(_customerMapper_expr) _customerMapper_select_clause.append(f"{_customerMapper_expr} AS latest_bill_id") _customerMapper_expr = """date_format(current_timestamp(), "yyyy-MM-dd\'T\'HH:mm:ss")""".replace("input_file_name()", "input_file") _customerMapper_expr = _customerMapper_expr.replace("_dq_source_file", "input_file") if "." in _customerMapper_expr: _customerMapper_expr = rewrite_response_body_json_access(_customerMapper_expr) _customerMapper_select_clause.append(f"{_customerMapper_expr} AS last_synced_at") _customerMapper_mapping_sql = ("SELECT " + ', '.join(_customerMapper_select_clause) + " FROM filter__1_df").replace("{job_id}", f"'{job_id}'") _customerMapper_input_data = { "component": "customerMapper", "datasource": "filter__1", "include_existing_columns": False, "to_schema_field_count": 7, } try: customerMapper_df = spark.sql(_customerMapper_mapping_sql) except AnalysisException as e: handle_analysis_error( e, component_name="customerMapper", error_code="TRF-MAP-002", exception_class=SchemaMappingException, message=f"Spark analysis error during customerMapper mapping: {e!s}", job_id=job_id, workspace=workspace, workflow=workflow, execution_environment=execution_environment, extra_details={"retry_job_id": retry_job_id or None, "sql_preview": _customerMapper_mapping_sql[:2000]}, input_data=_customerMapper_input_data, ) except Py4JJavaError as e: handle_java_error( e, component_name="customerMapper", operation="mapping SQL", format_name="sql", job_id=job_id, workspace=workspace, workflow=workflow, execution_environment=execution_environment, override_class=ExpressionException, override_code="TRF-EXP-001", extra_details={"retry_job_id": retry_job_id or None, "sql_preview": _customerMapper_mapping_sql[:2000]}, input_data=_customerMapper_input_data, ) customerMapper_df, customerMapper_observer = observe_metrics("customerMapper_df", customerMapper_df) customerMapper_df.createOrReplaceTempView("customerMapper_df") customerMapper_dependency_key="customerMapper" print(filter__1_dependency_key) customerMapper_execute_status="SUCCESS" except Exception as e: customerMapper_error = e log_error(LOGGER, f"Component customerMapper Failed", e, component_name="customerMapper") customerMapper_execute_status="ERROR" raise e finally: customerMapper_end_time=time.time() # %% customer_code_transform_start_time=time.time() try: # input dataframe will be connected component output as customerMapper_df # add processing logic here and create output as customer_code_transform_df customer_code_transform_df = spark.sql(""" SELECT * FROM ( SELECT *, ROW_NUMBER() OVER ( PARTITION BY account_id ORDER BY latest_bill_date DESC, latest_bill_id DESC ) AS rn FROM customerMapper_df ) t WHERE rn = 1 """) # TODO set output dataframe # --- Logging additions (safe from template-brace conflicts) --- output_df = customer_code_transform_df row_count = output_df.count() schema_str = output_df.schema.simpleString() print("Output row count:", row_count) print("Schema:", schema_str) output_df.show(5, truncate=False) customer_code_transform_df, customer_code_transform_observer = observe_metrics("customer_code_transform_df", customer_code_transform_df) customer_code_transform_df.createOrReplaceTempView("customer_code_transform_df") customer_code_transform_execute_status="SUCCESS" except Exception as e: print("ERROR:", str(e)) customer_code_transform_error = e log_error(LOGGER, f"Component customer_code_transform Failed", e) customer_code_transform_execute_status="ERROR" raise e customer_code_transform_end_time=time.time() customer_code_transform_dependency_key="customer_code_transform" print(customerMapper_dependency_key) # %% customerLatestBillMapper_start_time=time.time() customerLatestBillMapper_fail_on_error="True" try: _customerLatestBillMapper_select_clause=[] _customerLatestBillMapper_expr = """account_id""".replace("input_file_name()", "input_file") _customerLatestBillMapper_expr = _customerLatestBillMapper_expr.replace("_dq_source_file", "input_file") if "." in _customerLatestBillMapper_expr: _customerLatestBillMapper_expr = rewrite_response_body_json_access(_customerLatestBillMapper_expr) _customerLatestBillMapper_select_clause.append(f"{_customerLatestBillMapper_expr} AS account_id") _customerLatestBillMapper_expr = """id""".replace("input_file_name()", "input_file") _customerLatestBillMapper_expr = _customerLatestBillMapper_expr.replace("_dq_source_file", "input_file") if "." in _customerLatestBillMapper_expr: _customerLatestBillMapper_expr = rewrite_response_body_json_access(_customerLatestBillMapper_expr) _customerLatestBillMapper_select_clause.append(f"{_customerLatestBillMapper_expr} AS id") _customerLatestBillMapper_expr = """date_format(current_timestamp(), "yyyy-MM-dd\'T\'HH:mm:ss")""".replace("input_file_name()", "input_file") _customerLatestBillMapper_expr = _customerLatestBillMapper_expr.replace("_dq_source_file", "input_file") if "." in _customerLatestBillMapper_expr: _customerLatestBillMapper_expr = rewrite_response_body_json_access(_customerLatestBillMapper_expr) _customerLatestBillMapper_select_clause.append(f"{_customerLatestBillMapper_expr} AS last_synced_at") _customerLatestBillMapper_expr = """COALESCE(created_at, current_timestamp())""".replace("input_file_name()", "input_file") _customerLatestBillMapper_expr = _customerLatestBillMapper_expr.replace("_dq_source_file", "input_file") if "." in _customerLatestBillMapper_expr: _customerLatestBillMapper_expr = rewrite_response_body_json_access(_customerLatestBillMapper_expr) _customerLatestBillMapper_select_clause.append(f"{_customerLatestBillMapper_expr} AS created_at") _customerLatestBillMapper_expr = """status""".replace("input_file_name()", "input_file") _customerLatestBillMapper_expr = _customerLatestBillMapper_expr.replace("_dq_source_file", "input_file") if "." in _customerLatestBillMapper_expr: _customerLatestBillMapper_expr = rewrite_response_body_json_access(_customerLatestBillMapper_expr) _customerLatestBillMapper_select_clause.append(f"{_customerLatestBillMapper_expr} AS status") _customerLatestBillMapper_expr = """latest_bill_date""".replace("input_file_name()", "input_file") _customerLatestBillMapper_expr = _customerLatestBillMapper_expr.replace("_dq_source_file", "input_file") if "." in _customerLatestBillMapper_expr: _customerLatestBillMapper_expr = rewrite_response_body_json_access(_customerLatestBillMapper_expr) _customerLatestBillMapper_select_clause.append(f"{_customerLatestBillMapper_expr} AS latest_bill_date") _customerLatestBillMapper_expr = """latest_bill_id""".replace("input_file_name()", "input_file") _customerLatestBillMapper_expr = _customerLatestBillMapper_expr.replace("_dq_source_file", "input_file") if "." in _customerLatestBillMapper_expr: _customerLatestBillMapper_expr = rewrite_response_body_json_access(_customerLatestBillMapper_expr) _customerLatestBillMapper_select_clause.append(f"{_customerLatestBillMapper_expr} AS latest_bill_id") _customerLatestBillMapper_mapping_sql = ("SELECT " + ', '.join(_customerLatestBillMapper_select_clause) + " FROM customer_code_transform_df").replace("{job_id}", f"'{job_id}'") _customerLatestBillMapper_input_data = { "component": "customerLatestBillMapper", "datasource": "customer_code_transform", "include_existing_columns": False, "to_schema_field_count": 7, } try: customerLatestBillMapper_df = spark.sql(_customerLatestBillMapper_mapping_sql) except AnalysisException as e: handle_analysis_error( e, component_name="customerLatestBillMapper", error_code="TRF-MAP-002", exception_class=SchemaMappingException, message=f"Spark analysis error during customerLatestBillMapper mapping: {e!s}", job_id=job_id, workspace=workspace, workflow=workflow, execution_environment=execution_environment, extra_details={"retry_job_id": retry_job_id or None, "sql_preview": _customerLatestBillMapper_mapping_sql[:2000]}, input_data=_customerLatestBillMapper_input_data, ) except Py4JJavaError as e: handle_java_error( e, component_name="customerLatestBillMapper", operation="mapping SQL", format_name="sql", job_id=job_id, workspace=workspace, workflow=workflow, execution_environment=execution_environment, override_class=ExpressionException, override_code="TRF-EXP-001", extra_details={"retry_job_id": retry_job_id or None, "sql_preview": _customerLatestBillMapper_mapping_sql[:2000]}, input_data=_customerLatestBillMapper_input_data, ) customerLatestBillMapper_df, customerLatestBillMapper_observer = observe_metrics("customerLatestBillMapper_df", customerLatestBillMapper_df) customerLatestBillMapper_df.createOrReplaceTempView("customerLatestBillMapper_df") customerLatestBillMapper_dependency_key="customerLatestBillMapper" print(customer_code_transform_dependency_key) customerLatestBillMapper_execute_status="SUCCESS" except Exception as e: customerLatestBillMapper_error = e log_error(LOGGER, f"Component customerLatestBillMapper Failed", e, component_name="customerLatestBillMapper") customerLatestBillMapper_execute_status="ERROR" raise e finally: customerLatestBillMapper_end_time=time.time() # %% CheckpointOutput_start_time=time.time() CheckpointOutput_df = customerLatestBillMapper_df.localCheckpoint() customerLatestBillMapper_df.persist() CheckpointOutput_df.createOrReplaceTempView("CheckpointOutput_df") CheckpointOutput_end_time=time.time() CheckpointOutput_dependency_key="CheckpointOutput" print(customerLatestBillMapper_dependency_key) # %% customer_data_writer_start_time=time.time() customer_data_writer_fail_on_error="" try: _customer_data_writer_fields_to_update = CheckpointOutput_df.columns _customer_data_writer_set_clause=[] _customer_data_writer_unique_key_clause= [] for _key in ['account_id']: _customer_data_writer_unique_key_clause.append(f't.{_key} = s.{_key}') for _field in _customer_data_writer_fields_to_update: if(_field not in _customer_data_writer_unique_key_clause): _customer_data_writer_set_clause.append(f't.{_field} = s.{_field}') _merge_query = ''' MERGE INTO dremio.customer t USING CheckpointOutput_df s ON ''' + ' AND '.join(_customer_data_writer_unique_key_clause) + ''' WHEN MATCHED THEN UPDATE SET ''' + ', '.join(_customer_data_writer_set_clause) + ' WHEN NOT MATCHED THEN INSERT *' _customer_data_writer_input_data = {"component": "customer_data_writer", "datasource": "CheckpointOutput", "format": "iceberg", "mode": "merge"} try: spark.sql(_merge_query) except AnalysisException as e: handle_analysis_error( e, component_name="customer_data_writer", error_code="TRF-MRG-001", exception_class=MergeException, message=f"Merge failed for 'customer_data_writer' into dremio.customer: {e!s}", job_id=job_id, workspace=workspace, workflow=workflow, execution_environment=execution_environment, extra_details={"target_table": "dremio.customer", "merge_query_preview": _merge_query[:2000]}, input_data=_customer_data_writer_input_data, ) except Py4JJavaError as e: handle_java_error( e, component_name="customer_data_writer", operation="merge", format_name="iceberg", job_id=job_id, workspace=workspace, workflow=workflow, execution_environment=execution_environment, override_class=MergeException, override_code="TRF-MRG-001", extra_details={"target_table": "dremio.customer", "merge_query_preview": _merge_query[:2000]}, input_data=_customer_data_writer_input_data, ) customer_data_writer_dependency_key="customer_data_writer" print(CheckpointOutput_dependency_key) customer_data_writer_execute_status="SUCCESS" except Exception as e: customer_data_writer_error = e log_error(LOGGER, f"Component customer_data_writer Failed", e, component_name="customer_data_writer") customer_data_writer_execute_status="ERROR" raise e finally: customer_data_writer_end_time=time.time() # %% finalize_start_time=time.time() metrics = { 'data': collect_metrics(locals()), } materialization.materialized_execution_history({'finalize': {'execute_status': 'SUCCESS', 'fail_on_error': 'False', 'execution_order': os.environ.get('EXECUTION_ORDER')}, **metrics['data']}) log_info(LOGGER, f"Workflow Data metrics (correlation_id={metrics['data'].get('correlation_id')}): {metrics['data']}") finalize_end_time=time.time() if os.getenv('EXECUTION_ENVIRONMENT'): spark.stop()