{ "cells": [ { "cell_type": "code", "execution_count": 1, "id": "a286c8df", "metadata": {}, "outputs": [], "source": [ "# Copyright 2026 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": "d62dc22f", "metadata": {}, "source": [ "# Use AI Functions with the BigQuery Accessor\n", "\n", "\n", " \n", " \n", " \n", " \n", "
\n", " \n", " \"Colab Run in Colab\n", " \n", " \n", " \n", " \"GitHub\n", " View on GitHub\n", " \n", " \n", " \n", " \"Colab\n", " Open in Colab Enterprise\n", " \n", " \n", " \n", " \"BQ\n", " Open in BQ Studio\n", " \n", "
" ] }, { "cell_type": "markdown", "id": "fddd106d", "metadata": {}, "source": [ "## Environment Setup" ] }, { "cell_type": "markdown", "id": "53f8dd15", "metadata": {}, "source": [ "Make sure your GCP project have the follwing roles:\n", "* [roles/bigquery.jobUser](https://docs.cloud.google.com/iam/docs/roles-permissions/bigquery#bigquery.jobUser)\n", "* [roles/aiplatform.user](https://docs.cloud.google.com/iam/docs/roles-permissions/aiplatform#aiplatform.user)\n", "\n", "Then import bigframes to enable the `bigquery` accessor:" ] }, { "cell_type": "code", "execution_count": null, "id": "27900938", "metadata": {}, "outputs": [], "source": [ "import bigframes.pandas as bpd\n", "\n", "PROJECT_ID = \"\" # @param {type:\"string\"}\n", "LOCATION = \"US\" # @param {type:\"string\"}\n", "\n", "bpd.options.bigquery.project = PROJECT_ID\n", "bpd.options.bigquery.location = LOCATION\n", "bpd.options.display.progress_bar = None" ] }, { "cell_type": "markdown", "id": "d3e4540a", "metadata": {}, "source": [ "## Functions on pandas Series\n", "### Example: AI.EMBED" ] }, { "cell_type": "code", "execution_count": 3, "id": "b4f7f2f5", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "0 {'result': array([ 1.78243860e-03, -1.10658340...\n", "1 {'result': array([-7.29714287e-03, 1.04725976...\n", "dtype: struct, status: string>[pyarrow]" ] }, "metadata": {}, "output_type": "display_data" }, { "name": "stdout", "output_type": "stream", "text": [ "The type of the result is: \n" ] } ], "source": [ "import pandas as pd\n", "\n", "animals = pd.Series(['dog', 'fish'])\n", "result = animals.bigquery.ai.embed(endpoint='text-embedding-005')\n", "display(result)\n", "\n", "print(f\"The type of the result is: {type(result)}\")" ] }, { "cell_type": "markdown", "id": "3957ba3c", "metadata": {}, "source": [ "### Example: AI.SIMILARITY" ] }, { "cell_type": "markdown", "id": "1d021166", "metadata": {}, "source": [ "Computes similarities between a series and a constant:" ] }, { "cell_type": "code", "execution_count": 4, "id": "d62f63ee", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "0 0.629751\n", "1 0.768028\n", "dtype: Float64" ] }, "execution_count": 4, "metadata": {}, "output_type": "execute_result" } ], "source": [ "import pandas as pd\n", "\n", "animals = pd.Series(['dog', 'fish'])\n", "animals.bigquery.ai.similarity('shrimp', endpoint='text-embedding-005')" ] }, { "cell_type": "markdown", "id": "5918c440", "metadata": {}, "source": [ "Computes similiarities between two series:" ] }, { "cell_type": "code", "execution_count": 5, "id": "a1f61abf", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "0 0.651206\n", "1 0.835251\n", "dtype: Float64" ] }, "execution_count": 5, "metadata": {}, "output_type": "execute_result" } ], "source": [ "import pandas as pd\n", "\n", "animals = pd.Series(['dog', 'fish'])\n", "random_stuff = pd.Series(['smart phone', 'salmon'])\n", "\n", "animals.bigquery.ai.similarity(random_stuff, endpoint='text-embedding-005')\n" ] }, { "cell_type": "markdown", "id": "7df7a3d4", "metadata": {}, "source": [ "## Functions on pandas DataFrames\n", "### Example: AI.PREDICT\n", "\n", "First, prepare a dataset suitable for regression:" ] }, { "cell_type": "code", "execution_count": 6, "id": "a8b96156", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "=====Training data:\n" ] }, { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
neighborhoodbedroomsbathroomssqftmonthly_rent
0Downtown11.06502200
1Downtown21.59002800
2Downtown22.011003300
3Downtown32.515004200
4Suburbs21.59501600
5Suburbs32.013002100
6Suburbs32.516002500
7Suburbs43.021003100
8Midtown11.07001950
9Midtown21.59502500
10Midtown22.012003000
11Midtown32.516503800
\n", "
" ], "text/plain": [ " neighborhood bedrooms bathrooms sqft monthly_rent\n", "0 Downtown 1 1.0 650 2200\n", "1 Downtown 2 1.5 900 2800\n", "2 Downtown 2 2.0 1100 3300\n", "3 Downtown 3 2.5 1500 4200\n", "4 Suburbs 2 1.5 950 1600\n", "5 Suburbs 3 2.0 1300 2100\n", "6 Suburbs 3 2.5 1600 2500\n", "7 Suburbs 4 3.0 2100 3100\n", "8 Midtown 1 1.0 700 1950\n", "9 Midtown 2 1.5 950 2500\n", "10 Midtown 2 2.0 1200 3000\n", "11 Midtown 3 2.5 1650 3800" ] }, "metadata": {}, "output_type": "display_data" }, { "name": "stdout", "output_type": "stream", "text": [ "=====Prediction data:\n" ] }, { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
neighborhoodbedroomsbathroomssqft
0Downtown22.01050
1Suburbs32.01450
2Midtown11.5850
\n", "
" ], "text/plain": [ " neighborhood bedrooms bathrooms sqft\n", "0 Downtown 2 2.0 1050\n", "1 Suburbs 3 2.0 1450\n", "2 Midtown 1 1.5 850" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "import pandas as pd\n", "\n", "train_df = pd.DataFrame({\n", " \"neighborhood\": [\n", " \"Downtown\", \"Downtown\", \"Downtown\", \"Downtown\",\n", " \"Suburbs\", \"Suburbs\", \"Suburbs\", \"Suburbs\",\n", " \"Midtown\", \"Midtown\", \"Midtown\", \"Midtown\",\n", " ],\n", " \"bedrooms\": [1, 2, 2, 3, 2, 3, 3, 4, 1, 2, 2, 3],\n", " \"bathrooms\": [1.0, 1.5, 2.0, 2.5, 1.5, 2.0, 2.5, 3.0, 1.0, 1.5, 2.0, 2.5],\n", " \"sqft\": [650, 900, 1100, 1500, 950, 1300, 1600, 2100, 700, 950, 1200, 1650],\n", " \"monthly_rent\": [2200, 2800, 3300, 4200, 1600, 2100, 2500, 3100, 1950, 2500, 3000, 3800],\n", "})\n", "\n", "predict_df = pd.DataFrame({\n", " \"neighborhood\": [\"Downtown\", \"Suburbs\", \"Midtown\"],\n", " \"bedrooms\": [2, 3, 1],\n", " \"bathrooms\": [2.0, 2.0, 1.5],\n", " \"sqft\": [1050, 1450, 850],\n", "})\n", "\n", "print(\"=====Training data:\")\n", "display(train_df)\n", "print(\"=====Prediction data:\")\n", "display(predict_df)" ] }, { "cell_type": "markdown", "id": "ce41aabb", "metadata": {}, "source": [ "Then, perform a regression with **TabFM** without prior model training:" ] }, { "cell_type": "code", "execution_count": 7, "id": "0b93cf31", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
neighborhoodbedroomsbathroomssqftpredicted_monthly_rent
0Suburbs32.014502008.0
1Midtown11.58502352.0
2Downtown22.010503392.0
\n", "
" ], "text/plain": [ " neighborhood bedrooms bathrooms sqft predicted_monthly_rent\n", "0 Suburbs 3 2.0 1450 2008.0\n", "1 Midtown 1 1.5 850 2352.0\n", "2 Downtown 2 2.0 1050 3392.0" ] }, "execution_count": 7, "metadata": {}, "output_type": "execute_result" } ], "source": [ "predictions = train_df.bigquery.ai.predict(\n", " predict_df,\n", " label_col=\"monthly_rent\",\n", ")\n", "predictions" ] }, { "cell_type": "markdown", "id": "f3689dcb", "metadata": {}, "source": [ "### Example: AI.FORECAST\n", "\n", "First, prepare a timeseries dataset:" ] }, { "cell_type": "code", "execution_count": 8, "id": "b2588d49", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
dateitem_idsales
02025-01-01item_1100
12025-01-02item_1105
22025-01-03item_1110
32025-01-04item_1115
42025-01-05item_1120
52025-01-06item_1125
62025-01-07item_1130
72025-01-08item_1135
82025-01-09item_1140
92025-01-10item_1145
102025-01-11item_1150
112025-01-12item_1155
122025-01-13item_1160
132025-01-14item_1165
142025-01-01item_250
152025-01-02item_252
162025-01-03item_255
172025-01-04item_258
182025-01-05item_260
192025-01-06item_262
202025-01-07item_265
212025-01-08item_268
222025-01-09item_270
232025-01-10item_272
242025-01-11item_275
252025-01-12item_278
262025-01-13item_280
272025-01-14item_285
\n", "
" ], "text/plain": [ " date item_id sales\n", "0 2025-01-01 item_1 100\n", "1 2025-01-02 item_1 105\n", "2 2025-01-03 item_1 110\n", "3 2025-01-04 item_1 115\n", "4 2025-01-05 item_1 120\n", "5 2025-01-06 item_1 125\n", "6 2025-01-07 item_1 130\n", "7 2025-01-08 item_1 135\n", "8 2025-01-09 item_1 140\n", "9 2025-01-10 item_1 145\n", "10 2025-01-11 item_1 150\n", "11 2025-01-12 item_1 155\n", "12 2025-01-13 item_1 160\n", "13 2025-01-14 item_1 165\n", "14 2025-01-01 item_2 50\n", "15 2025-01-02 item_2 52\n", "16 2025-01-03 item_2 55\n", "17 2025-01-04 item_2 58\n", "18 2025-01-05 item_2 60\n", "19 2025-01-06 item_2 62\n", "20 2025-01-07 item_2 65\n", "21 2025-01-08 item_2 68\n", "22 2025-01-09 item_2 70\n", "23 2025-01-10 item_2 72\n", "24 2025-01-11 item_2 75\n", "25 2025-01-12 item_2 78\n", "26 2025-01-13 item_2 80\n", "27 2025-01-14 item_2 85" ] }, "execution_count": 8, "metadata": {}, "output_type": "execute_result" } ], "source": [ "import pandas as pd\n", "\n", "dates = pd.date_range(\"2025-01-01\", periods=14, freq=\"D\")\n", "\n", "df = pd.DataFrame({\n", " \"date\": dates.tolist() * 2,\n", " \"item_id\": [\"item_1\"] * 14 + [\"item_2\"] * 14,\n", " \"sales\": [\n", " 100, 105, 110, 115, 120, 125, 130, 135, 140, 145, 150, 155, 160, 165,\n", " 50, 52, 55, 58, 60, 62, 65, 68, 70, 72, 75, 78, 80, 85,\n", " ],\n", "})\n", "\n", "df\n" ] }, { "cell_type": "markdown", "id": "28cdfb02", "metadata": {}, "source": [ "Then, use TimesFM to forecase the time series" ] }, { "cell_type": "code", "execution_count": 9, "id": "419e2ebe", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
item_idforecast_timestampforecast_valueconfidence_levelprediction_interval_lower_boundprediction_interval_upper_boundai_forecast_status
0item_12025-01-16 00:00:00+00:00174.2574920.95175.728216177.85742
1item_12025-01-15 00:00:00+00:00167.4660030.95169.34789169.553343
2item_12025-01-17 00:00:00+00:00180.4913330.95178.587954185.040092
3item_22025-01-17 00:00:00+00:0091.997520.9579.145413100.786089
4item_22025-01-15 00:00:00+00:0087.9501040.9582.87669892.470649
5item_22025-01-16 00:00:00+00:0090.0503460.9581.46053396.463373
\n", "
" ], "text/plain": [ " item_id forecast_timestamp forecast_value confidence_level \\\n", "0 item_1 2025-01-16 00:00:00+00:00 174.257492 0.95 \n", "1 item_1 2025-01-15 00:00:00+00:00 167.466003 0.95 \n", "2 item_1 2025-01-17 00:00:00+00:00 180.491333 0.95 \n", "3 item_2 2025-01-17 00:00:00+00:00 91.99752 0.95 \n", "4 item_2 2025-01-15 00:00:00+00:00 87.950104 0.95 \n", "5 item_2 2025-01-16 00:00:00+00:00 90.050346 0.95 \n", "\n", " prediction_interval_lower_bound prediction_interval_upper_bound \\\n", "0 175.728216 177.85742 \n", "1 169.34789 169.553343 \n", "2 178.587954 185.040092 \n", "3 79.145413 100.786089 \n", "4 82.876698 92.470649 \n", "5 81.460533 96.463373 \n", "\n", " ai_forecast_status \n", "0 \n", "1 \n", "2 \n", "3 \n", "4 \n", "5 " ] }, "execution_count": 9, "metadata": {}, "output_type": "execute_result" } ], "source": [ "df.bigquery.ai.forecast(\n", " data_col=\"sales\", \n", " timestamp_col=\"date\", \n", " id_cols=[\"item_id\"], \n", " horizon=3\n", ")\n" ] }, { "cell_type": "markdown", "id": "2ffa0534", "metadata": {}, "source": [ "### Example: AI.GENERATE_BOOL\n", "Evaluates a structured prompt condition using Gemini and returns a boolean series:" ] }, { "cell_type": "code", "execution_count": 10, "id": "eba4b82f", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "0 True\n", "1 False\n", "Name: result, dtype: bool[pyarrow]" ] }, "execution_count": 10, "metadata": {}, "output_type": "execute_result" } ], "source": [ "import pandas as pd\n", "\n", "df = pd.DataFrame({\n", " \"animal\": [\"cougar\", \"fish\"],\n", " \"habitat\": [\"mountain\", \"grassland\"]\n", "})\n", "\n", "prompt = (df[\"animal\"], \"lives in the \", df[\"habitat\"])\n", "\n", "df.bigquery.ai.generate_bool(prompt, endpoint=\"gemini-2.5-flash\").struct.field(\"result\")\n" ] } ], "metadata": { "kernelspec": { "display_name": "venv (3.14.2)", "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.14.2" } }, "nbformat": 4, "nbformat_minor": 5 }