{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "new_license_cell_0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright 2025 Google LLC\n",
    "#\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "#     https://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b2300e6",
   "metadata": {},
   "source": [
    "# Asset Price Forecast using Iceberg and Prophet"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9440e57a",
   "metadata": {},
   "source": [
    "<img src=\"../../docs/images/forecast/gold-forecast.png\" alt=\"drawing\" width=\"800\"/>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "new_overview_2",
   "metadata": {},
   "source": [
    "## Overview\n",
    "This notebook demonstrates how to build an asset price forecasting solution using Dataproc Serverless Spark Connect with Iceberg tables and the Prophet forecasting library. It covers data loading, preprocessing, model training, prediction, and storing results back into Iceberg."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "new_setup_env_3",
   "metadata": {},
   "source": [
    "## Setup Environment"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "new_install_libs_4",
   "metadata": {},
   "source": [
    "### Install Required Libraries\n",
    "This cell installs the necessary Python libraries for this notebook, including `pyspark`, `google-spark-connect`, `google-cloud-dataproc`, `pandas`, `prophet`, and `matplotlib`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6892eb2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -q dataproc-spark-connect pandas prophet matplotlib"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "new_import_libs_5",
   "metadata": {},
   "source": [
    "### Import Libraries\n",
    "Import all the required libraries for this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d06431b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from prophet import Prophet\n",
    "\n",
    "from pyspark.sql.functions import col, year\n",
    "\n",
    "from google.cloud.dataproc_spark_connect import DataprocSparkSession"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "new_configure_gcp_6",
   "metadata": {},
   "source": [
    "## Configure Google Cloud Project and Resources"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "new_set_config_7",
   "metadata": {},
   "source": [
    "### Set Configuration Variables\n",
    "Define the Google Cloud project ID, region, Dataproc Serverless template ID, and GCS paths for input data and the Iceberg warehouse. Replace the placeholder values with your actual project details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8860965",
   "metadata": {},
   "outputs": [],
   "source": [
    "project_id = \"YOUR_PROJECT_ID\"\n",
    "location = \"YOUR_LOCATION\"\n",
    "\n",
    "csv_path = 'gs://dataproc-metastore-public-binaries/asset_price_forecast/asset_price_forecast.csv'\n",
    "\n",
    "iceberg_warehouse_gcs_path = \"gs://your-bucket-iceberg-warehouse/iceberg/data\"\n",
    "iceberg_catalog_name = \"biglake\"\n",
    "iceberg_dataset = \"finance\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "new_create_template_8",
   "metadata": {},
   "source": [
    "### Create Dataproc Serverless Spark Connect Session\n",
    "This function creates a Dataproc Serverless Spark Session. This session configured to use Iceberg with BigQuery as the metastore, enabling seamless integration between Spark and Iceberg tables stored in Google Cloud Storage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de001e63",
   "metadata": {},
   "outputs": [],
   "source": [
    "spark = (\n",
    "    DataprocSparkSession.builder\n",
    "    .projectId(project_id)\n",
    "    .location(location)\n",
    "    .config(f\"spark.sql.catalog.{iceberg_catalog_name}\", \"org.apache.iceberg.spark.SparkCatalog\")\n",
    "    .config(f\"spark.sql.catalog.{iceberg_catalog_name}.type\", \"hadoop\")\n",
    "    .config(f\"spark.sql.catalog.{iceberg_catalog_name}.warehouse\", iceberg_warehouse_gcs_path)\n",
    "    .getOrCreate()\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "new_data_loading_10",
   "metadata": {},
   "source": [
    "## Data Loading and Preparation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "new_load_data_11",
   "metadata": {},
   "source": [
    "### Load Raw Data from GCS\n",
    "Load the asset price data from a CSV file stored in Google Cloud Storage into a Spark DataFrame. The `header` option is set to `true` to indicate that the first row is a header, and `inferSchema` is set to `true` to automatically detect column types."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c6797c4-e1be-4eca-9775-12e86128b81c",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = spark.read.option(\"header\", \"true\").option(\"inferSchema\", \"true\").csv(csv_path)\n",
    "df.count()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "new_preprocess_data_12",
   "metadata": {},
   "source": [
    "### Preprocess Data\n",
    "Perform basic data preprocessing by dropping unnecessary columns ('High', 'Low', 'Open', 'Volume') and extracting the year from the 'date' column. This optimized DataFrame will be used for forecasting and storage in Iceberg."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b25618aa-f819-424b-a34b-695aa848f4a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_optimized = df.drop('High', 'Low', 'Open', 'Volume')\n",
    "df_optimized = df_optimized.withColumn('year', year(col('date')))\n",
    "df_optimized.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "new_store_iceberg_13",
   "metadata": {},
   "source": [
    "## Store Data in Iceberg Table"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "new_create_iceberg_db_14",
   "metadata": {},
   "source": [
    "### Create Iceberg Database\n",
    "Create a new database within your Iceberg catalog. This logical container will hold your Iceberg tables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bba976f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "spark.sql(f\"CREATE DATABASE IF NOT EXISTS {iceberg_catalog_name}.{iceberg_dataset}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "new_write_optimized_15",
   "metadata": {},
   "source": [
    "### Write Optimized Data to Iceberg Table\n",
    "Write the preprocessed Spark DataFrame into an Iceberg table. The table is partitioned by 'year' for optimized querying, and the `overwrite` mode ensures that any existing table with the same name is replaced."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88a78a2c-3209-4356-92c9-a548d8f6c92b",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_optimized.write.format('iceberg').mode('overwrite').partitionBy('year').saveAsTable(f'{iceberg_catalog_name}.{iceberg_dataset}.gold_price')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "new_prophet_forecast_16",
   "metadata": {},
   "source": [
    "## Time Series Forecasting with Prophet"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "new_spark_to_pandas_17",
   "metadata": {},
   "source": [
    "### Convert Spark DataFrame to Pandas DataFrame\n",
    "Convert the Spark DataFrame to a Pandas DataFrame, which is required by the Prophet library for time series modeling."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92137c35-0e0c-49bb-874a-b1bf5ab95e2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_pandas = df_optimized.toPandas()\n",
    "df_pandas.count()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "new_train_prophet_18",
   "metadata": {},
   "source": [
    "### Train the Prophet Model\n",
    "Rename the 'Date' column to 'ds' and 'Close' column to 'y' to match Prophet's expected input format. Initialize and train the Prophet model with a specified confidence interval."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6daf8726-c67c-4c50-98a5-64cff93a2a14",
   "metadata": {},
   "outputs": [],
   "source": [
    "series = df_pandas.rename(columns={'Date': 'ds', 'Close': 'y'})\n",
    "confidence_interval = 0.9\n",
    "model = Prophet(interval_width=confidence_interval)\n",
    "model.fit(series)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "new_generate_future_19",
   "metadata": {},
   "source": [
    "### Generate Future Dates and Make Predictions\n",
    "Generate a DataFrame with future dates for the next 365 days and use the trained Prophet model to make predictions for these future dates, along with upper and lower bounds for the forecast."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "294ee85a-ee13-4700-a00e-986087f8926b",
   "metadata": {},
   "outputs": [],
   "source": [
    "forecast_period = 365\n",
    "future = model.make_future_dataframe(periods=forecast_period)\n",
    "forecast = model.predict(future)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "new_visualize_forecast_20",
   "metadata": {},
   "source": [
    "### Visualize the Forecast\n",
    "Plot the historical data and the forecasted values, including the confidence interval. The plot is customized with a title and axis labels for clarity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b3f253f-40c7-4208-9c7a-f4987b60e85c",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig1 = model.plot(forecast)\n",
    "fig1.gca().set_title(\"Gold Spot Price Forecast\", size=16)\n",
    "fig1.gca().set_xlabel(\"Date\")\n",
    "fig1.gca().set_ylabel(\"Price USD/Ounce\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "new_store_forecasted_21",
   "metadata": {},
   "source": [
    "## Store Forecasted Data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "new_prepare_forecast_spark_22",
   "metadata": {},
   "source": [
    "### Prepare Forecast Data for Spark\n",
    "Filter the forecast to include only future predictions, rename columns to match the original DataFrame, convert the 'Date' column type, and then convert the Pandas DataFrame back into a Spark DataFrame for consistency and further operations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eea8fe86-0554-495c-9da2-df82a8dc1e55",
   "metadata": {},
   "outputs": [],
   "source": [
    "forecast1 = forecast[['ds', 'yhat']][forecast['ds']> '2025-07-17']\n",
    "forecast1.columns = ['Date', 'Close']\n",
    "forecast1['Date'] = forecast1['Date'].dt.date\n",
    "df_forecast = spark.createDataFrame(forecast1)\n",
    "df_forecast = df_forecast.withColumn('year', year(col('date')))\n",
    "df_forecast.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "new_combine_data_23",
   "metadata": {},
   "source": [
    "### Combine Historical and Forecasted Data\n",
    "Combine the original optimized historical data with the newly generated forecasted data into a single Spark DataFrame."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09c226ca-53a2-48c5-885c-83adba602be9",
   "metadata": {},
   "outputs": [],
   "source": [
    "combined_df = df_optimized.unionByName(df_forecast)\n",
    "combined_df.tail(20)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "new_append_iceberg_24",
   "metadata": {},
   "source": [
    "### Append Combined Data to Iceberg Table\n",
    "Append the combined historical and forecasted data to the existing Iceberg table. This demonstrates how to update your Iceberg table with new data incrementally."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6626681-753a-4162-83cb-61a0c5099108",
   "metadata": {},
   "outputs": [],
   "source": [
    "combined_df.write.format('iceberg').mode('append').partitionBy('year').saveAsTable(f'{iceberg_catalog_name}.{iceberg_dataset}.gold_price')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}