{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "outputs": [],
      "source": [
        "# Copyright 2023 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",
      "metadata": {
        "id": "JAPoU8Sm5E6e"
      },
      "source": [
        "# Get started with BigQuery DataFrames\n",
        "\n",
        "<table align=\"left\">\n",
        "\n",
        "  <td>\n",
        "    <a href=\"https://colab.research.google.com/github/googleapis/python-bigquery-dataframes/blob/main/notebooks/getting_started/bq_dataframes_template.ipynb\">\n",
        "      <img src=\"https://raw.githubusercontent.com/googleapis/python-bigquery-dataframes/refs/heads/main/third_party/logo/colab-logo.png\" alt=\"Colab logo\"> Run in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://github.com/googleapis/python-bigquery-dataframes/blob/main/notebooks/getting_started/bq_dataframes_template.ipynb\">\n",
        "      <img src=\"https://raw.githubusercontent.com/googleapis/python-bigquery-dataframes/refs/heads/main/third_party/logo/github-logo.png\" width=\"32\" alt=\"GitHub logo\">\n",
        "      View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/googleapis/python-bigquery-dataframes/blob/main/notebooks/getting_started/bq_dataframes_template.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
        "      Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://console.cloud.google.com/bigquery/import?url=https://github.com/googleapis/python-bigquery-dataframes/blob/main/notebooks/getting_started/bq_dataframes_template.ipynb\">\n",
        "      <img src=\"https://encrypted-tbn0.gstatic.com/images?q=tbn:ANd9GcTW1gvOovVlbZAIZylUtf5Iu8-693qS1w5NJw&s\" alt=\"BQ logo\" width=\"35\">\n",
        "      Open in BQ Studio\n",
        "    </a>\n",
        "  </td>\n",
        "\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "24743cf4a1e1"
      },
      "source": [
        "**_NOTE_**: This notebook has been tested in the following environment:\n",
        "\n",
        "* Python version = 3.10"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "BigQuery DataFrames (also known as BigFrames) provides a Pythonic DataFrame and machine learning (ML) API powered by the BigQuery engine.\n",
        "\n",
        "* `bigframes.pandas` provides a pandas API for analytics. Many workloads can be\n",
        "  migrated from pandas to bigframes by just changing a few imports.\n",
        "* `bigframes.ml` provides a scikit-learn-like API for ML.\n",
        "* `bigframes.ml.llm` provides API for large language models including Gemini.\n",
        "\n",
        "You can learn more about [BigQuery DataFrames](https://cloud.google.com/bigquery/docs/bigquery-dataframes-introduction) and its [API reference](https://cloud.google.com/python/docs/reference/bigframes/latest).\n",
        "\n",
        "For any issues or feedback please reach out to bigframes-feedback@google.com."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BF1j6f9HApxa"
      },
      "source": [
        "## Before you begin\n",
        "\n",
        "Complete the tasks in this section to set up your environment."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Yq7zKYWelRQP"
      },
      "source": [
        "### Install the python package\n",
        "\n",
        "You need the [bigframes](https://pypi.org/project/bigframes/) python package to be installed. If you don't have that, uncomment and run the following cell and *restart the kernel*."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [],
      "source": [
        "#%pip install  --upgrade bigframes"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WReHDGG5g0XY"
      },
      "source": [
        "### Set your project id and location\n",
        "\n",
        "Following are some quick references:\n",
        "\n",
        "* Google Cloud Project: https://cloud.google.com/resource-manager/docs/creating-managing-projects.\n",
        "* BigQuery Location: https://cloud.google.com/bigquery/docs/locations."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "oM1iC_MfAts1"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"bigframes-dev\"  # @param {type: \"string\"}\n",
        "LOCATION = \"US\"  # @param {type: \"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "960505627ddf"
      },
      "source": [
        "### Import library"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "PyQmSRbKA8r-"
      },
      "outputs": [],
      "source": [
        "import bigframes.pandas as bpd"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "init_aip:mbsdk,all"
      },
      "source": [
        "\n",
        "### Set BigQuery DataFrames options"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "NPPMuw2PXGeo"
      },
      "outputs": [],
      "source": [
        "# Note: The project option is not required in all environments.\n",
        "# For example, In BigQuery Studio, the project ID is automatically detected,\n",
        "# But in Google Colab it must be set by the user.\n",
        "bpd.options.bigquery.project = PROJECT_ID\n",
        "\n",
        "# Note: The location option is not required.\n",
        "# It defaults to the location of the first table or query\n",
        "# passed to read_gbq(). For APIs where a location can't be\n",
        "# auto-detected, the location defaults to the \"US\" location.\n",
        "bpd.options.bigquery.location = LOCATION\n",
        "\n",
        "# Note: BigQuery DataFrames objects are by default fully ordered like Pandas.\n",
        "# If ordering is not important for you, you can uncomment the following\n",
        "# expression to run BigQuery DataFrames in partial ordering mode.\n",
        "bpd.options.bigquery.ordering_mode = \"partial\"\n",
        "\n",
        "# Note: By default BigQuery DataFrames emits out BigQuery job metadata via a\n",
        "# progress bar. But in this notebook let's disable the progress bar to keep the\n",
        "# experience less verbose. If you would like the default behavior, please\n",
        "# comment out the following expression. \n",
        "bpd.options.display.progress_bar = None"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pDfrKwMKE_dK"
      },
      "source": [
        "If you want to reset the project and/or location of the created DataFrame or Series objects, reset the session by executing `bpd.close_session()`. After that, you can redo the above steps."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9EMAqR37AfLS"
      },
      "source": [
        "## Create a BigQuery DataFrames DataFrame\n",
        "\n",
        "You can create a BigQuery DataFrames DataFrame by reading data from any of the\n",
        "following locations:\n",
        "\n",
        "* A local data file\n",
        "* Data stored in a BigQuery table\n",
        "* A data file stored in Cloud Storage\n",
        "* An in-memory pandas DataFrame\n",
        "\n",
        "Note that the DataFrame does not copy the data to the local memory, instead\n",
        "keeps the underlying data in a BigQuery table during read and analysis. That's\n",
        "how it can handle really large size of data (at BigQuery Scale) independent of\n",
        "the local memory.\n",
        "\n",
        "For simplicity, speed and cost efficiency, this tutorial uses the\n",
        "[`penguins`](https://pantheon.corp.google.com/bigquery?ws=!1m5!1m4!4m3!1sbigquery-public-data!2sml_datasets!3spenguins)\n",
        "table from BigQuery public data, which contains 27 KB data about a set of\n",
        "penguins - species, island of residence, culmen length and depth, flipper length\n",
        "and sex. There is a version of this data in the Cloud Storage\n",
        "[cloud samples data](https://pantheon.corp.google.com/storage/browser/_details/cloud-samples-data/vertex-ai/bigframe/penguins.csv)\n",
        "as well."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "Vyex9BQI-BNa"
      },
      "outputs": [],
      "source": [
        "# This is how you read a BigQuery table\n",
        "df = bpd.read_gbq(\"bigquery-public-data.ml_datasets.penguins\")\n",
        "\n",
        "# This is how you would read a csv from the Cloud Storage\n",
        "#df = bpd.read_csv(\"gs://cloud-samples-data/vertex-ai/bigframe/penguins.csv\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We can use `peek` to preview a few rows (selected arbitrarily) from the dataframes:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>species</th>\n",
              "      <th>island</th>\n",
              "      <th>culmen_length_mm</th>\n",
              "      <th>culmen_depth_mm</th>\n",
              "      <th>flipper_length_mm</th>\n",
              "      <th>body_mass_g</th>\n",
              "      <th>sex</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>Adelie Penguin (Pygoscelis adeliae)</td>\n",
              "      <td>Dream</td>\n",
              "      <td>36.6</td>\n",
              "      <td>18.4</td>\n",
              "      <td>184.0</td>\n",
              "      <td>3475.0</td>\n",
              "      <td>FEMALE</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>Adelie Penguin (Pygoscelis adeliae)</td>\n",
              "      <td>Dream</td>\n",
              "      <td>39.8</td>\n",
              "      <td>19.1</td>\n",
              "      <td>184.0</td>\n",
              "      <td>4650.0</td>\n",
              "      <td>MALE</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2</th>\n",
              "      <td>Adelie Penguin (Pygoscelis adeliae)</td>\n",
              "      <td>Dream</td>\n",
              "      <td>40.9</td>\n",
              "      <td>18.9</td>\n",
              "      <td>184.0</td>\n",
              "      <td>3900.0</td>\n",
              "      <td>MALE</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3</th>\n",
              "      <td>Chinstrap penguin (Pygoscelis antarctica)</td>\n",
              "      <td>Dream</td>\n",
              "      <td>46.5</td>\n",
              "      <td>17.9</td>\n",
              "      <td>192.0</td>\n",
              "      <td>3500.0</td>\n",
              "      <td>FEMALE</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>4</th>\n",
              "      <td>Adelie Penguin (Pygoscelis adeliae)</td>\n",
              "      <td>Dream</td>\n",
              "      <td>37.3</td>\n",
              "      <td>16.8</td>\n",
              "      <td>192.0</td>\n",
              "      <td>3000.0</td>\n",
              "      <td>FEMALE</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>"
            ],
            "text/plain": [
              "                                     species island  culmen_length_mm  \\\n",
              "0        Adelie Penguin (Pygoscelis adeliae)  Dream              36.6   \n",
              "1        Adelie Penguin (Pygoscelis adeliae)  Dream              39.8   \n",
              "2        Adelie Penguin (Pygoscelis adeliae)  Dream              40.9   \n",
              "3  Chinstrap penguin (Pygoscelis antarctica)  Dream              46.5   \n",
              "4        Adelie Penguin (Pygoscelis adeliae)  Dream              37.3   \n",
              "\n",
              "   culmen_depth_mm  flipper_length_mm  body_mass_g     sex  \n",
              "0             18.4              184.0       3475.0  FEMALE  \n",
              "1             19.1              184.0       4650.0    MALE  \n",
              "2             18.9              184.0       3900.0    MALE  \n",
              "3             17.9              192.0       3500.0  FEMALE  \n",
              "4             16.8              192.0       3000.0  FEMALE  "
            ]
          },
          "execution_count": 7,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "df.peek()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gE6CEALjDZZV"
      },
      "source": [
        "We just created a DataFrame, `df`, refering to the entirety of the source table data, without downloading it to the local machine."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rwPLjqW2Ajzh"
      },
      "source": [
        "## Inspect and manipulate data in BigQuery DataFrames"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bExmYlL_ELtV"
      },
      "source": [
        "### Using pandas API\n",
        "\n",
        "You can use pandas API on the BigQuery DataFrames DataFrame as you normally would in Pandas, but computation happens in the BigQuery query engine instead of your local environment."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EJIZJaNXFQzh"
      },
      "source": [
        "Let's compute the mean of the `body_mass_g` series:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "YKwCW7Nsavap"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "average_body_mass: 4201.754385964914\n"
          ]
        }
      ],
      "source": [
        "average_body_mass = df[\"body_mass_g\"].mean()\n",
        "print(f\"average_body_mass: {average_body_mass}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DSs1cnca-MOU"
      },
      "source": [
        "Calculate the mean `body_mass_g` by `species` using the `groupby` operation:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "4PyKMR61-Mjy"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>body_mass_g</th>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>species</th>\n",
              "      <th></th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>Adelie Penguin (Pygoscelis adeliae)</th>\n",
              "      <td>3700.662252</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>Chinstrap penguin (Pygoscelis antarctica)</th>\n",
              "      <td>3733.088235</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>Gentoo penguin (Pygoscelis papua)</th>\n",
              "      <td>5076.01626</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "<p>3 rows × 1 columns</p>\n",
              "</div>[3 rows x 1 columns in total]"
            ],
            "text/plain": [
              "                                           body_mass_g\n",
              "species                                               \n",
              "Adelie Penguin (Pygoscelis adeliae)        3700.662252\n",
              "Chinstrap penguin (Pygoscelis antarctica)  3733.088235\n",
              "Gentoo penguin (Pygoscelis papua)           5076.01626\n",
              "\n",
              "[3 rows x 1 columns]"
            ]
          },
          "execution_count": 9,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "df[[\"species\", \"body_mass_g\"]].groupby(by=df[\"species\"]).mean(numeric_only=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6sf9kZ2C9Ixe"
      },
      "source": [
        "You can confirm that the calculations were run in BigQuery by clicking \"Open job\" from the previous cells' output. This takes you to the BigQuery console to view the SQL statement and job details."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Using SQL functions\n",
        "\n",
        "The [bigframes.bigquery module](https://cloud.google.com/python/docs/reference/bigframes/latest/bigframes.bigquery) provides many [BigQuery SQL functions](https://cloud.google.com/bigquery/docs/reference/standard-sql/functions-all) which may not have a pandas-equivalent."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {},
      "outputs": [],
      "source": [
        "import bigframes.bigquery"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The `bigframes.bigquery.struct()` function creates a new STRUCT Series with subfields for each column in a DataFrames."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "0    {'culmen_length_mm': 36.6, 'culmen_depth_mm': ...\n",
              "1    {'culmen_length_mm': 39.8, 'culmen_depth_mm': ...\n",
              "2    {'culmen_length_mm': 40.9, 'culmen_depth_mm': ...\n",
              "3    {'culmen_length_mm': 46.5, 'culmen_depth_mm': ...\n",
              "4    {'culmen_length_mm': 37.3, 'culmen_depth_mm': ...\n",
              "dtype: struct<culmen_length_mm: double, culmen_depth_mm: double, flipper_length_mm: double>[pyarrow]"
            ]
          },
          "execution_count": 11,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "lengths = bigframes.bigquery.struct(\n",
        "    df[[\"culmen_length_mm\", \"culmen_depth_mm\", \"flipper_length_mm\"]]\n",
        ")\n",
        "lengths.peek()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Use the `bigframes.bigquery.sql_scalar()` function to access arbitrary SQL syntax representing a single column expression."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "0    18.4\n",
              "1    19.1\n",
              "2    18.9\n",
              "3    17.9\n",
              "4    16.8\n",
              "dtype: Float64"
            ]
          },
          "execution_count": 12,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "shortest = bigframes.bigquery.sql_scalar(\n",
        "    \"LEAST({0}, {1}, {2})\",\n",
        "    columns=[df['culmen_depth_mm'], df['culmen_length_mm'], df['flipper_length_mm']],\n",
        ")\n",
        "shortest.peek()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Visualize data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### First party visualizations\n",
        "\n",
        "BigQuery DataFrames provides a number of visualizations via the `plot` method and [accessor](https://cloud.google.com/python/docs/reference/bigframes/latest/bigframes.operations.plotting.PlotAccessor) on the DataFrame and Series objects."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "<Axes: title={'center': 'Numeric features'}>"
            ]
          },
          "execution_count": 13,
          "metadata": {},
          "output_type": "execute_result"
        },
        {
          "data": {
            "image/png": "[encoded data, 67892 characters]",
            "text/plain": [
              "<Figure size 640x480 with 1 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "(\n",
        "    df\n",
        "    .sort_values(\"body_mass_g\")\n",
        "    .reset_index(drop=True)\n",
        "    .plot(title=\"Numeric features\", logy=True)\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "<Axes: xlabel='species'>"
            ]
          },
          "execution_count": 14,
          "metadata": {},
          "output_type": "execute_result"
        },
        {
          "data": {
            "image/png": "[encoded data, 45412 characters]",
            "text/plain": [
              "<Figure size 640x480 with 1 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "means = df.groupby(\"species\").mean(numeric_only=True)\n",
        "means.plot.bar()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Integration with open source visualizations\n",
        "\n",
        "BigQuery Dataframes is also compatible with several open source visualization packages, such as `matplotlib`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "image/png": "[encoded data, 36604 characters]",
            "text/plain": [
              "<Figure size 640x480 with 1 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "\n",
        "# plotting a histogram\n",
        "species_counts = df[\"species\"].value_counts()\n",
        "plt.pie(species_counts, labels=species_counts.index, autopct='%1.1f%%')\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Pandas interoperability\n",
        "\n",
        "BigQuery DataFrames can be converted from and to Pandas DataFrame with `to_pandas` and `read_pandas` respectively.\n",
        "This could be handy to take advantage of the capabilities of the two systems.\n",
        "\n",
        "> Note: `to_pandas` converts the BigQuery DataFrame to Pandas DataFrame by bringing all the data in memory, which would be an issue\n",
        "for large data, as your machine may not have enough memory to accommodate that."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "We have a dataframe of <class 'bigframes.dataframe.DataFrame'>\n",
            "\n",
            "\n",
            "We have a dataframe of <class 'pandas.core.frame.DataFrame'>\n",
            "\n",
            "\n",
            "We have a dataframe of <class 'bigframes.dataframe.DataFrame'>\n",
            "\n"
          ]
        }
      ],
      "source": [
        "def print_type(df):\n",
        "    print(f\"\\nWe have a dataframe of {type(df)}\\n\")\n",
        "\n",
        "# The original bigframes dataframe\n",
        "cur_df = df\n",
        "print_type(cur_df)\n",
        "\n",
        "# Convert to pandas dataframe\n",
        "cur_df = cur_df.to_pandas()\n",
        "print_type(cur_df)\n",
        "\n",
        "# Convert back to bigframes dataframe\n",
        "cur_df = bpd.read_pandas(cur_df)\n",
        "print_type(cur_df)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Machine Learning with BigQuery DataFrames"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Clean and prepare data\n",
        "\n",
        "We're are going to start with supervised learning, where a Linear Regression model will learn to predict the body mass (output variable `y`) using input features such as flipper length, sex, species, and more (features `X`)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Drop any rows that has missing (NA) values\n",
        "df = df.dropna()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Part of preparing data for a machine learning task is splitting it into subsets for training and testing to ensure that the solution is not overfitting. By default, BQML will automatically manage splitting the data for you. However, BQML also supports manually splitting out your training data.\n",
        "\n",
        "Performing a manual data split can be done with `bigframes.ml.model_selection.train_test_split` like so:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "    df_train shape: (267, 7)\n",
            "    df_test shape: (67, 7)\n",
            "\n"
          ]
        }
      ],
      "source": [
        "from bigframes.ml.model_selection import train_test_split\n",
        "\n",
        "\n",
        "# This will split df into test and training sets, with 20% of the rows in the test set,\n",
        "# and the rest in the training set\n",
        "df_train, df_test = train_test_split(df, test_size=0.2)\n",
        "\n",
        "# Show the shape of the data after the split\n",
        "print(f\"\"\"\n",
        "    df_train shape: {df_train.shape}\n",
        "    df_test shape: {df_test.shape}\n",
        "\"\"\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "    X_train shape: (267, 6)\n",
            "    X_test shape: (67, 7)\n",
            "    y_train shape: (267, 1)\n",
            "    y_test shape: (67, 1)\n",
            "\n"
          ]
        }
      ],
      "source": [
        "# Isolate input features and output variable into DataFrames\n",
        "X_train = df_train[[\n",
        "    'island',\n",
        "    'culmen_length_mm',\n",
        "    'culmen_depth_mm',\n",
        "    'flipper_length_mm',\n",
        "    'sex',\n",
        "    'species',\n",
        "]]\n",
        "y_train = df_train[['body_mass_g']]\n",
        "\n",
        "X_test = df_test[[\n",
        "    'island',\n",
        "    'culmen_length_mm',\n",
        "    'culmen_depth_mm',\n",
        "    'flipper_length_mm',\n",
        "    'sex',\n",
        "    'species',\n",
        "    # Include the actual body_mass_g so that we can compare with the predicted\n",
        "    # without a join.\n",
        "    'body_mass_g'\n",
        "]]\n",
        "y_test = df_test[['body_mass_g']]\n",
        "\n",
        "# Print the shapes of features and label\n",
        "print(f\"\"\"\n",
        "    X_train shape: {X_train.shape}\n",
        "    X_test shape: {X_test.shape}\n",
        "    y_train shape: {y_train.shape}\n",
        "    y_test shape: {y_test.shape}\n",
        "\"\"\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Define pipeline\n",
        "\n",
        "This step is subjective to the problem. Although a model can be directly trained on the original data, it is often useful to apply some preprocessing to the original data.\n",
        "In this example we want to apply a [`ColumnTransformer`](https://cloud.google.com/python/docs/reference/bigframes/latest/bigframes.ml.compose.ColumnTransformer) in which we apply [`OneHotEncoder`](https://cloud.google.com/python/docs/reference/bigframes/latest/bigframes.ml.preprocessing.OneHotEncoder) to the category features and [`StandardScaler`](https://cloud.google.com/python/docs/reference/bigframes/latest/bigframes.ml.preprocessing.StandardScaler) to the numeric features."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Pipeline(steps=[('preproc',\n",
              "                 ColumnTransformer(transformers=[('onehot', OneHotEncoder(),\n",
              "                                                  ['island', 'species', 'sex']),\n",
              "                                                 ('scaler', StandardScaler(),\n",
              "                                                  ['culmen_depth_mm',\n",
              "                                                   'culmen_length_mm',\n",
              "                                                   'flipper_length_mm'])])),\n",
              "                ('linreg', LinearRegression(fit_intercept=False))])"
            ]
          },
          "execution_count": 20,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "from bigframes.ml.linear_model import LinearRegression\n",
        "from bigframes.ml.pipeline import Pipeline\n",
        "from bigframes.ml.compose import ColumnTransformer\n",
        "from bigframes.ml.preprocessing import StandardScaler, OneHotEncoder\n",
        "\n",
        "preprocessing = ColumnTransformer([\n",
        "    (\"onehot\", OneHotEncoder(), [\"island\", \"species\", \"sex\"]),\n",
        "    (\"scaler\", StandardScaler(), [\"culmen_depth_mm\", \"culmen_length_mm\", \"flipper_length_mm\"]),\n",
        "])\n",
        "\n",
        "model = LinearRegression(fit_intercept=False)\n",
        "\n",
        "pipeline = Pipeline([\n",
        "    ('preproc', preprocessing),\n",
        "    ('linreg', model)\n",
        "])\n",
        "\n",
        "# View the pipeline\n",
        "pipeline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Train and Predict\n",
        "\n",
        "Supervised learning is when we train a model on input-output pairs, and then ask it to predict the output for new inputs. An example of such a predictor is `bigframes.ml.linear_models.LinearRegression`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>predicted_body_mass_g</th>\n",
              "      <th>island</th>\n",
              "      <th>culmen_length_mm</th>\n",
              "      <th>culmen_depth_mm</th>\n",
              "      <th>flipper_length_mm</th>\n",
              "      <th>sex</th>\n",
              "      <th>species</th>\n",
              "      <th>body_mass_g</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>3271.548077</td>\n",
              "      <td>Biscoe</td>\n",
              "      <td>37.9</td>\n",
              "      <td>18.6</td>\n",
              "      <td>172.0</td>\n",
              "      <td>FEMALE</td>\n",
              "      <td>Adelie Penguin (Pygoscelis adeliae)</td>\n",
              "      <td>3150.0</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>3224.661209</td>\n",
              "      <td>Biscoe</td>\n",
              "      <td>37.7</td>\n",
              "      <td>16.0</td>\n",
              "      <td>183.0</td>\n",
              "      <td>FEMALE</td>\n",
              "      <td>Adelie Penguin (Pygoscelis adeliae)</td>\n",
              "      <td>3075.0</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2</th>\n",
              "      <td>3395.403541</td>\n",
              "      <td>Biscoe</td>\n",
              "      <td>34.5</td>\n",
              "      <td>18.1</td>\n",
              "      <td>187.0</td>\n",
              "      <td>FEMALE</td>\n",
              "      <td>Adelie Penguin (Pygoscelis adeliae)</td>\n",
              "      <td>2900.0</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3</th>\n",
              "      <td>3943.436439</td>\n",
              "      <td>Biscoe</td>\n",
              "      <td>40.1</td>\n",
              "      <td>18.9</td>\n",
              "      <td>188.0</td>\n",
              "      <td>MALE</td>\n",
              "      <td>Adelie Penguin (Pygoscelis adeliae)</td>\n",
              "      <td>4300.0</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>4</th>\n",
              "      <td>3986.662895</td>\n",
              "      <td>Biscoe</td>\n",
              "      <td>41.4</td>\n",
              "      <td>18.6</td>\n",
              "      <td>191.0</td>\n",
              "      <td>MALE</td>\n",
              "      <td>Adelie Penguin (Pygoscelis adeliae)</td>\n",
              "      <td>3700.0</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>"
            ],
            "text/plain": [
              "   predicted_body_mass_g  island  culmen_length_mm  culmen_depth_mm  \\\n",
              "0            3271.548077  Biscoe              37.9             18.6   \n",
              "1            3224.661209  Biscoe              37.7             16.0   \n",
              "2            3395.403541  Biscoe              34.5             18.1   \n",
              "3            3943.436439  Biscoe              40.1             18.9   \n",
              "4            3986.662895  Biscoe              41.4             18.6   \n",
              "\n",
              "   flipper_length_mm     sex                              species  body_mass_g  \n",
              "0              172.0  FEMALE  Adelie Penguin (Pygoscelis adeliae)       3150.0  \n",
              "1              183.0  FEMALE  Adelie Penguin (Pygoscelis adeliae)       3075.0  \n",
              "2              187.0  FEMALE  Adelie Penguin (Pygoscelis adeliae)       2900.0  \n",
              "3              188.0    MALE  Adelie Penguin (Pygoscelis adeliae)       4300.0  \n",
              "4              191.0    MALE  Adelie Penguin (Pygoscelis adeliae)       3700.0  "
            ]
          },
          "execution_count": 21,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# Learn from the training data how to predict output y\n",
        "pipeline.fit(X_train, y_train)\n",
        "\n",
        "# Predict y for the test data\n",
        "y_pred = pipeline.predict(X_test)\n",
        "\n",
        "# View predictions preview\n",
        "y_pred.peek()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Evaluate results\n",
        "\n",
        "Some models include a convenient `.score(X, y)` method for evaulation with a preset accuracy metric:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>mean_absolute_error</th>\n",
              "      <th>mean_squared_error</th>\n",
              "      <th>mean_squared_log_error</th>\n",
              "      <th>median_absolute_error</th>\n",
              "      <th>r2_score</th>\n",
              "      <th>explained_variance</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>231.914252</td>\n",
              "      <td>78873.600421</td>\n",
              "      <td>0.005172</td>\n",
              "      <td>178.724985</td>\n",
              "      <td>0.890549</td>\n",
              "      <td>0.890566</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "<p>1 rows × 6 columns</p>\n",
              "</div>[1 rows x 6 columns in total]"
            ],
            "text/plain": [
              " mean_absolute_error  mean_squared_error  mean_squared_log_error  \\\n",
              "          231.914252        78873.600421                0.005172   \n",
              "\n",
              " median_absolute_error  r2_score  explained_variance  \n",
              "            178.724985  0.890549            0.890566  \n",
              "\n",
              "[1 rows x 6 columns]"
            ]
          },
          "execution_count": 22,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "pipeline.score(X_test.drop(columns=[\"body_mass_g\"]), y_test)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "For a more general approach, the library `bigframes.ml.metrics` is provided:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 23,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "np.float64(0.8905492944632485)"
            ]
          },
          "execution_count": 23,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "from bigframes.ml.metrics import r2_score\n",
        "\n",
        "r2_score(y_pred['body_mass_g'], y_pred[\"predicted_body_mass_g\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Generative AI with BigQuery DataFrames\n",
        "\n",
        "BigQuery DataFrames integration with the Large Language Models (LLM) supported by BigQuery ML. Check out the [`bigframes.ml.llm`](https://cloud.google.com/python/docs/reference/bigframes/latest/bigframes.ml.llm) module for all the available models.\n",
        "\n",
        "To use this feature you would need to have a few additional APIs enabled and IAM roles configured. Please make sure of that by following [this documentation](https://cloud.google.com/bigquery/docs/use-bigquery-dataframes#remote-models) and then uncomment the code in the following cells to try out the integration with Gemini."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Create prompts\n",
        "\n",
        "A \"prompt\" text column can be initialized either directly or via the pandas APIs. For simplicity let's use a direct initialization here."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>prompt</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>What is BigQuery?</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>What is BQML?</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2</th>\n",
              "      <td>What is BigQuery DataFrames?</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "<p>3 rows × 1 columns</p>\n",
              "</div>[3 rows x 1 columns in total]"
            ],
            "text/plain": [
              "                         prompt\n",
              "0             What is BigQuery?\n",
              "1                 What is BQML?\n",
              "2  What is BigQuery DataFrames?\n",
              "\n",
              "[3 rows x 1 columns]"
            ]
          },
          "execution_count": 24,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "df = bpd.DataFrame(\n",
        "        {\n",
        "            \"prompt\": [\"What is BigQuery?\", \"What is BQML?\", \"What is BigQuery DataFrames?\"],\n",
        "        })\n",
        "df"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Generate responses\n",
        "\n",
        "Here we will use the [`GeminiTextGenerator`](https://cloud.google.com/python/docs/reference/bigframes/latest/bigframes.ml.llm.GeminiTextGenerator) LLM to answer the questions. Read the [GeminiTextGenerator API documentation](https://cloud.google.com/python/docs/reference/bigframes/latest/bigframes.ml.llm.GeminiTextGenerator) for all the model versions supported via the `model_name` param."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "metadata": {},
      "outputs": [],
      "source": [
        "# from bigframes.ml.llm import GeminiTextGenerator\n",
        "\n",
        "# model = GeminiTextGenerator(model_name=\"gemini-2.0-flash-001\")\n",
        "\n",
        "# pred = model.predict(df)\n",
        "# pred"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's print the full text response for the question \"What is BigQuery DataFrames?\"."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 26,
      "metadata": {},
      "outputs": [],
      "source": [
        "# import IPython.display\n",
        "\n",
        "# IPython.display.Markdown(pred.loc[2][\"ml_generate_text_llm_result\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TpV-iwP9qw9c"
      },
      "source": [
        "## Cleaning up\n",
        "\n",
        "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
        "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
        "\n",
        "To remove any temporary cloud artifacts (inclusing BQ tables) created in the current BigQuery DataFrames session, simply call `close_session`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 27,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Delete the temporary cloud artifacts created during the bigframes session \n",
        "bpd.close_session()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wCsmt0IwFkDy"
      },
      "source": [
        "## Summary and next steps\n",
        "\n",
        "1. You created BigQuery DataFrames objects, and inspected and manipulated data with pandas APIs at BigQuery scale and speed.\n",
        "\n",
        "1. You also created ML model from a DataFrame and used them to run predictions on another DataFrame.\n",
        "\n",
        "1. You got access to Google's state-of-the-art Gemini LLM through simple pythonic API.\n",
        "\n",
        "Learn more about BigQuery DataFrames in the documentation [BigQuery DataFrames](https://cloud.google.com/bigquery/docs/bigquery-dataframes-introduction) and its [API reference](https://cloud.google.com/python/docs/reference/bigframes/latest).\n",
        "\n",
        "Also, find more sample notebooks in the [GitHub repo](https://github.com/googleapis/python-bigquery-dataframes/tree/main/notebooks), including the [pypi.ipynb](https://github.com/googleapis/python-bigquery-dataframes/blob/main/notebooks/dataframes/pypi.ipynb) that processes 400+ TB data at the cost and efficiency close to direct SQL by taking advantage of the [partial ordering](https://cloud.google.com/python/docs/reference/bigframes/latest/bigframes._config.bigquery_options.BigQueryOptions#bigframes__config_bigquery_options_BigQueryOptions_ordering_mode) mode."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "toc_visible": true
    },
    "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.10.16"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
