diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS new file mode 100644 index 00000000000..ef110646417 --- /dev/null +++ b/.github/CODEOWNERS @@ -0,0 +1,12 @@ +# Code owners file. +# This file controls who is tagged for review for any given pull request. +# +# For syntax help see: +# https://help.github.com/en/github/creating-cloning-and-archiving-repositories/about-code-owners#codeowners-syntax +# Note: This file is autogenerated. To make changes to the codeowner team, please update .repo-metadata.json. + +# @googleapis/cloud-sdk-python-team @googleapis/bigquery-dataframe-team are the default owners for changes in this repo +* @googleapis/cloud-sdk-python-team @googleapis/bigquery-dataframe-team + +# @googleapis/python-samples-reviewers @googleapis/bigquery-dataframe-team are the default owners for samples changes +/samples/ @googleapis/python-samples-reviewers @googleapis/bigquery-dataframe-team diff --git a/.github/CONTRIBUTING.md b/.github/CONTRIBUTING.md new file mode 100644 index 00000000000..939e5341e74 --- /dev/null +++ b/.github/CONTRIBUTING.md @@ -0,0 +1,28 @@ +# How to Contribute + +We'd love to accept your patches and contributions to this project. There are +just a few small guidelines you need to follow. + +## Contributor License Agreement + +Contributions to this project must be accompanied by a Contributor License +Agreement. You (or your employer) retain the copyright to your contribution; +this simply gives us permission to use and redistribute your contributions as +part of the project. Head over to to see +your current agreements on file or to sign a new one. + +You generally only need to submit a CLA once, so if you've already submitted one +(even if it was for a different project), you probably don't need to do it +again. + +## Code reviews + +All submissions, including submissions by project members, require review. We +use GitHub pull requests for this purpose. Consult +[GitHub Help](https://help.github.com/articles/about-pull-requests/) for more +information on using pull requests. + +## Community Guidelines + +This project follows [Google's Open Source Community +Guidelines](https://opensource.google.com/conduct/). diff --git a/.github/ISSUE_TEMPLATE/bug_report.md b/.github/ISSUE_TEMPLATE/bug_report.md new file mode 100644 index 00000000000..0745497ddf2 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/bug_report.md @@ -0,0 +1,58 @@ +--- +name: Bug report +about: Create a report to help us improve + +--- + +Thanks for stopping by to let us know something could be better! + +**PLEASE READ**: If you have a support contract with Google, please create an issue in the [support console](https://cloud.google.com/support/) instead of filing on GitHub. This will ensure a timely response. + +Please run down the following list and make sure you've tried the usual "quick fixes": + + - Search the issues already opened: https://github.com/googleapis/python-bigquery-dataframes/issues + - Search StackOverflow: https://stackoverflow.com/questions/tagged/google-cloud-platform+python + +If you are still having issues, please be sure to include as much information as possible: + +#### Environment details + + - OS type and version: + - Python version: `python --version` + - pip version: `pip --version` + - `bigframes` version: `pip show bigframes` + + +```python +import sys +import bigframes +import google.cloud.bigquery +import pandas +import pyarrow + +print(f"Python: {sys.version}") +print(f"bigframes=={bigframes.__version__}") +print(f"google-cloud-bigquery=={google.cloud.bigquery.__version__}") +print(f"pandas=={pandas.__version__}") +print(f"pyarrow=={pyarrow.__version__}") +``` + +#### Steps to reproduce + + 1. ? + 2. ? + +#### Code example + +```python +# example +``` + +#### Stack trace +``` +# example +``` + +Making sure to follow these steps will guarantee the quickest resolution possible. + +Thanks! diff --git a/.github/ISSUE_TEMPLATE/feature_request.md b/.github/ISSUE_TEMPLATE/feature_request.md new file mode 100644 index 00000000000..6365857f33c --- /dev/null +++ b/.github/ISSUE_TEMPLATE/feature_request.md @@ -0,0 +1,18 @@ +--- +name: Feature request +about: Suggest an idea for this library + +--- + +Thanks for stopping by to let us know something could be better! + +**PLEASE READ**: If you have a support contract with Google, please create an issue in the [support console](https://cloud.google.com/support/) instead of filing on GitHub. This will ensure a timely response. + + **Is your feature request related to a problem? Please describe.** +A clear and concise description of what the problem is. Ex. I'm always frustrated when [...] + **Describe the solution you'd like** +A clear and concise description of what you want to happen. + **Describe alternatives you've considered** +A clear and concise description of any alternative solutions or features you've considered. + **Additional context** +Add any other context or screenshots about the feature request here. diff --git a/.github/ISSUE_TEMPLATE/support_request.md b/.github/ISSUE_TEMPLATE/support_request.md new file mode 100644 index 00000000000..99586903212 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/support_request.md @@ -0,0 +1,7 @@ +--- +name: Support request +about: If you have a support contract with Google, please create an issue in the Google Cloud Support console. + +--- + +**PLEASE READ**: If you have a support contract with Google, please create an issue in the [support console](https://cloud.google.com/support/) instead of filing on GitHub. This will ensure a timely response. diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md new file mode 100644 index 00000000000..3e59d9a70d1 --- /dev/null +++ b/.github/PULL_REQUEST_TEMPLATE.md @@ -0,0 +1,7 @@ +Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly: +- [ ] Make sure to open an issue as a [bug/issue](https://github.com/googleapis/python-bigquery-dataframes/issues/new/choose) before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea +- [ ] Ensure the tests and linter pass +- [ ] Code coverage does not decrease (if any source code was changed) +- [ ] Appropriate docs were updated (if necessary) + +Fixes # 🦕 diff --git a/.github/auto-label.yaml b/.github/auto-label.yaml new file mode 100644 index 00000000000..21786a4eb08 --- /dev/null +++ b/.github/auto-label.yaml @@ -0,0 +1,20 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +requestsize: + enabled: true + +path: + pullrequest: true + paths: + samples: "samples" diff --git a/.github/blunderbuss.yml b/.github/blunderbuss.yml new file mode 100644 index 00000000000..527a5e8f9a0 --- /dev/null +++ b/.github/blunderbuss.yml @@ -0,0 +1,17 @@ +# Blunderbuss config +# +# This file controls who is assigned for pull requests and issues. +# Note: This file is autogenerated. To make changes to the assignee +# team, please update `codeowner_team` in `.repo-metadata.json`. +assign_issues: + - googleapis/bigquery-dataframe-team + +assign_issues_by: + - labels: + - "samples" + to: + - googleapis/python-samples-reviewers + - googleapis/bigquery-dataframe-team + +assign_prs: + - googleapis/bigquery-dataframe-team diff --git a/.github/header-checker-lint.yml b/.github/header-checker-lint.yml new file mode 100644 index 00000000000..6fe78aa7987 --- /dev/null +++ b/.github/header-checker-lint.yml @@ -0,0 +1,15 @@ +{"allowedCopyrightHolders": ["Google LLC"], + "allowedLicenses": ["Apache-2.0", "MIT", "BSD-3"], + "ignoreFiles": ["**/requirements.txt", "**/requirements-test.txt", "**/__init__.py", "samples/**/constraints.txt", "samples/**/constraints-test.txt"], + "sourceFileExtensions": [ + "ts", + "js", + "java", + "sh", + "Dockerfile", + "yaml", + "py", + "html", + "txt" + ] +} \ No newline at end of file diff --git a/.github/snippet-bot.yml b/.github/snippet-bot.yml new file mode 100644 index 00000000000..e69de29bb2d diff --git a/.github/workflows/docs-deploy.yml b/.github/workflows/docs-deploy.yml new file mode 100644 index 00000000000..13d4d87263b --- /dev/null +++ b/.github/workflows/docs-deploy.yml @@ -0,0 +1,57 @@ +name: Deploy docs to GitHub Pages + +on: + # Runs on pushes targeting the default branch + # TODO(tswast): Update this to only be releases once we confirm it's working. + push: + branches: ["main"] + + # Allows you to run this workflow manually from the Actions tab + workflow_dispatch: + +# Sets permissions of the GITHUB_TOKEN to allow deployment to GitHub Pages +permissions: + contents: read + pages: write + id-token: write + +# Allow only one concurrent deployment, skipping runs queued between the run in-progress and latest queued. +# However, do NOT cancel in-progress runs as we want to allow these production deployments to complete. +concurrency: + group: "pages" + cancel-in-progress: false + +jobs: + # Build job + build: + runs-on: ubuntu-latest + steps: + - name: Checkout + uses: actions/checkout@v4 + - name: Setup Python + uses: actions/setup-python@v5 + with: + python-version: "3.13" + - name: Install nox + run: | + python -m pip install --upgrade setuptools pip wheel + python -m pip install nox + - name: Run docs + run: | + nox -s docs + - name: Upload artifact + uses: actions/upload-pages-artifact@v3 + with: + path: docs/_build/html/ + + # Deployment job + deploy: + environment: + name: github-pages + url: ${{ steps.deployment.outputs.page_url }} + runs-on: ubuntu-latest + needs: build + steps: + - name: Deploy to GitHub Pages + id: deployment + uses: actions/deploy-pages@v4 diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml new file mode 100644 index 00000000000..cf8fe9cce25 --- /dev/null +++ b/.github/workflows/docs.yml @@ -0,0 +1,38 @@ +on: + pull_request: + branches: + - main + push: + branches: + - main +name: docs +jobs: + docs: + runs-on: ubuntu-latest + steps: + - name: Checkout + uses: actions/checkout@v4 + - name: Setup Python + uses: actions/setup-python@v5 + with: + python-version: "3.13" + - name: Install nox + run: | + python -m pip install --upgrade setuptools pip wheel + python -m pip install nox + - name: Run docs + run: | + nox -s docs + docfx: + runs-on: ubuntu-latest + steps: + - name: Checkout + uses: actions/checkout@v4 + - name: Setup Python + uses: actions/setup-python@v5 + with: + python-version: "3.10" + - name: Install nox + run: | + python -m pip install --upgrade setuptools pip wheel + python -m pip install nox diff --git a/.github/workflows/js-tests.yml b/.github/workflows/js-tests.yml new file mode 100644 index 00000000000..588aa854f32 --- /dev/null +++ b/.github/workflows/js-tests.yml @@ -0,0 +1,20 @@ +name: js-tests +on: + pull_request: + branches: + - main + push: + branches: + - main +jobs: + build: + runs-on: ubuntu-latest + steps: + - name: Checkout + uses: actions/checkout@v4 + - name: Install modules + working-directory: ./tests/js + run: npm install + - name: Run tests + working-directory: ./tests/js + run: npm test diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml new file mode 100644 index 00000000000..b848262c3aa --- /dev/null +++ b/.github/workflows/lint.yml @@ -0,0 +1,28 @@ +on: + pull_request: + branches: + - main + push: + branches: + - main +name: lint +jobs: + lint: + runs-on: ubuntu-latest + steps: + - name: Checkout + uses: actions/checkout@v4 + - name: Setup Python + uses: actions/setup-python@v5 + with: + python-version: "3.14" + - name: Install nox + run: | + python -m pip install --upgrade setuptools pip wheel + python -m pip install nox + - name: Run lint + run: | + nox -s lint + - name: Run lint_setup_py + run: | + nox -s lint_setup_py diff --git a/.github/workflows/mypy.yml b/.github/workflows/mypy.yml new file mode 100644 index 00000000000..fc9e9709461 --- /dev/null +++ b/.github/workflows/mypy.yml @@ -0,0 +1,25 @@ +on: + pull_request: + branches: + - main + push: + branches: + - main +name: mypy +jobs: + mypy: + runs-on: ubuntu-latest + steps: + - name: Checkout + uses: actions/checkout@v4 + - name: Setup Python + uses: actions/setup-python@v5 + with: + python-version: "3.10" + - name: Install nox + run: | + python -m pip install --upgrade setuptools pip wheel + python -m pip install nox + - name: Run mypy + run: | + nox -s mypy diff --git a/.github/workflows/unittest.yml b/.github/workflows/unittest.yml new file mode 100644 index 00000000000..f3e3dae6c5e --- /dev/null +++ b/.github/workflows/unittest.yml @@ -0,0 +1,64 @@ +on: + pull_request: + branches: + - main + push: + branches: + - main +name: unittest +jobs: + unit: + # TODO(https://github.com/googleapis/gapic-generator-python/issues/2303): use `ubuntu-latest` once this bug is fixed. + # Use ubuntu-22.04 until Python 3.7 is removed from the test matrix + # https://docs.github.com/en/actions/using-github-hosted-runners/using-github-hosted-runners/about-github-hosted-runners#standard-github-hosted-runners-for-public-repositories + runs-on: ubuntu-22.04 + strategy: + matrix: + python: ['3.10', '3.11', '3.12', '3.13', '3.14'] + steps: + - name: Checkout + uses: actions/checkout@v4 + - name: Setup Python + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python }} + - name: Install nox + run: | + python -m pip install --upgrade setuptools pip wheel + python -m pip install nox + - name: Run unit tests + env: + COVERAGE_FILE: .coverage-${{ matrix.python }} + run: | + nox -s unit-${{ matrix.python }} + - name: Upload coverage results + uses: actions/upload-artifact@v4 + with: + name: coverage-artifact-${{ matrix.python }} + path: .coverage-${{ matrix.python }} + include-hidden-files: true + + cover: + runs-on: ubuntu-latest + needs: + - unit + steps: + - name: Checkout + uses: actions/checkout@v4 + - name: Setup Python + uses: actions/setup-python@v5 + with: + python-version: "3.14" + - name: Install coverage + run: | + python -m pip install --upgrade setuptools pip wheel + python -m pip install coverage + - name: Download coverage results + uses: actions/download-artifact@v4 + with: + path: .coverage-results/ + - name: Report coverage results + run: | + find .coverage-results -type f -name '*.zip' -exec unzip {} \; + coverage combine .coverage-results/**/.coverage* + coverage report --show-missing --fail-under=35 diff --git a/.gitignore b/.gitignore new file mode 100644 index 00000000000..6b157559ccd --- /dev/null +++ b/.gitignore @@ -0,0 +1,70 @@ +*.py[cod] +*.sw[op] + +# C extensions +*.so + +# Packages +*.egg +*.egg-info +dist +build +eggs +.eggs +parts +bin +var +sdist +develop-eggs +.installed.cfg +lib +lib64 +__pycache__ + +# Installer logs +pip-log.txt + +# Unit test / coverage reports +.coverage +.nox +.cache +.pytest_cache + + +# Mac +.DS_Store + +# JetBrains +.idea + +# VS Code +.vscode + +# emacs +*~ + +# Built documentation +docs/_build +bigquery/docs/generated +docs.metadata + +# Virtual environment +env/ +venv/ + +# Test logs +coverage.xml +*sponge_log.xml + +# System test environment variables. +system_tests/local_test_setup +tests/js/node_modules/ + +# Make sure a generated file isn't accidentally committed. +pylintrc +pylintrc.test +dummy.pkl +.mypy_cache/ + +# Gemini +GEMINI.md diff --git a/.kokoro/build.sh b/.kokoro/build.sh new file mode 100755 index 00000000000..6cc03455da4 --- /dev/null +++ b/.kokoro/build.sh @@ -0,0 +1,55 @@ +#!/bin/bash +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +set -eo pipefail + +PROJECT_SCM="github/python-bigquery-dataframes" + +if [[ -z "${PROJECT_ROOT:-}" ]]; then + PROJECT_ROOT="${KOKORO_ARTIFACTS_DIR}/${PROJECT_SCM}" +fi + +cd "${PROJECT_ROOT}" + +# Disable buffering, so that the logs stream through. +export PYTHONUNBUFFERED=1 + +# Workaround https://github.com/pytest-dev/pytest/issues/9567 +export PY_IGNORE_IMPORTMISMATCH=1 + +# Debug: show build environment +env | grep KOKORO + +# Install pip +python3 -m pip install --upgrade --quiet pip +python3 -m pip --version + +# Remove old nox +python3 -m pip uninstall --yes --quiet nox-automation + +# Install nox +python3 -m pip install --upgrade --quiet nox +python3 -m nox --version + +# If NOX_SESSION is set, it only runs the specified session, +# otherwise run all the sessions. +if [[ -n "${NOX_SESSION:-}" ]]; then + python3 -m nox --stop-on-first-error -s ${NOX_SESSION:-} +else + python3 -m nox --stop-on-first-error +fi + +# Prevent kokoro from trying to collect many mb of artifacts, wasting several minutes +sudo rm -rf "${KOKORO_ARTIFACTS_DIR?}"/* diff --git a/.kokoro/continuous/common.cfg b/.kokoro/continuous/common.cfg new file mode 100644 index 00000000000..97e0651aa92 --- /dev/null +++ b/.kokoro/continuous/common.cfg @@ -0,0 +1,10 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +# Build logs will be here +action { + define_artifacts { + regex: "**/*sponge_log.xml" + } +} + +build_file: "python-bigquery-dataframes/.kokoro/build.sh" diff --git a/.kokoro/continuous/continuous.cfg b/.kokoro/continuous/continuous.cfg new file mode 100644 index 00000000000..8f43917d92f --- /dev/null +++ b/.kokoro/continuous/continuous.cfg @@ -0,0 +1 @@ +# Format: //devtools/kokoro/config/proto/build.proto \ No newline at end of file diff --git a/.kokoro/continuous/doctest.cfg b/.kokoro/continuous/doctest.cfg new file mode 100644 index 00000000000..2aad95beed4 --- /dev/null +++ b/.kokoro/continuous/doctest.cfg @@ -0,0 +1,12 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +# Only run this nox session. +env_vars: { + key: "NOX_SESSION" + value: "cleanup doctest" +} + +env_vars: { + key: "GOOGLE_CLOUD_PROJECT" + value: "bigframes-testing" +} diff --git a/.kokoro/continuous/e2e.cfg b/.kokoro/continuous/e2e.cfg new file mode 100644 index 00000000000..e049dd30b3a --- /dev/null +++ b/.kokoro/continuous/e2e.cfg @@ -0,0 +1,17 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +# Only run this nox session. +env_vars: { + key: "NOX_SESSION" + value: "e2e unit_prerelease system_prerelease system_noextras" +} + +env_vars: { + key: "GOOGLE_CLOUD_PROJECT" + value: "bigframes-load-testing" +} + +env_vars: { + key: "BIGFRAMES_TEST_MODEL_VERTEX_ENDPOINT" + value: "https://us-central1-aiplatform.googleapis.com/v1/projects/272725758477/locations/us-central1/endpoints/590545496255234048" +} diff --git a/.kokoro/continuous/nightly.cfg b/.kokoro/continuous/nightly.cfg new file mode 100644 index 00000000000..2b7111664f7 --- /dev/null +++ b/.kokoro/continuous/nightly.cfg @@ -0,0 +1,3 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +build_file: "python-bigquery-dataframes/.kokoro/release-nightly.sh" diff --git a/.kokoro/continuous/notebook.cfg b/.kokoro/continuous/notebook.cfg new file mode 100644 index 00000000000..cc73c3bea40 --- /dev/null +++ b/.kokoro/continuous/notebook.cfg @@ -0,0 +1,12 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +# Only run this nox session. +env_vars: { + key: "NOX_SESSION" + value: "notebook" +} + +env_vars: { + key: "GOOGLE_CLOUD_PROJECT" + value: "bigframes-testing" +} diff --git a/.kokoro/continuous/prerelease-deps.cfg b/.kokoro/continuous/prerelease-deps.cfg new file mode 100644 index 00000000000..3595fb43f5c --- /dev/null +++ b/.kokoro/continuous/prerelease-deps.cfg @@ -0,0 +1,7 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +# Only run this nox session. +env_vars: { + key: "NOX_SESSION" + value: "prerelease_deps" +} diff --git a/.kokoro/continuous/windows.cfg b/.kokoro/continuous/windows.cfg new file mode 100644 index 00000000000..806986138dd --- /dev/null +++ b/.kokoro/continuous/windows.cfg @@ -0,0 +1,3 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +build_file: "python-bigquery-dataframes/scripts/windows/build.bat" diff --git a/.kokoro/load/benchmark.cfg b/.kokoro/load/benchmark.cfg new file mode 100644 index 00000000000..45e2527a0c1 --- /dev/null +++ b/.kokoro/load/benchmark.cfg @@ -0,0 +1,22 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +# Only run this nox session. +env_vars: { + key: "NOX_SESSION" + value: "benchmark" +} + +env_vars: { + key: "BENCHMARK_AND_PUBLISH" + value: "true" +} + +env_vars: { + key: "GOOGLE_CLOUD_PROJECT" + value: "bigframes-benchmarking" +} + +env_vars: { + key: "BIGFRAMES_TEST_MODEL_VERTEX_ENDPOINT" + value: "https://us-central1-aiplatform.googleapis.com/v1/projects/272725758477/locations/us-central1/endpoints/590545496255234048" +} diff --git a/.kokoro/load/common.cfg b/.kokoro/load/common.cfg new file mode 100644 index 00000000000..d86932662d4 --- /dev/null +++ b/.kokoro/load/common.cfg @@ -0,0 +1,11 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +# Build logs will be here +action { + define_artifacts { + regex: "**/*sponge_log.xml" + } +} + +build_file: "python-bigquery-dataframes/.kokoro/build.sh" +timeout_mins: 720 diff --git a/.kokoro/load/load.cfg b/.kokoro/load/load.cfg new file mode 100644 index 00000000000..656614cf735 --- /dev/null +++ b/.kokoro/load/load.cfg @@ -0,0 +1,17 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +# Only run this nox session. +env_vars: { + key: "NOX_SESSION" + value: "load" +} + +env_vars: { + key: "GOOGLE_CLOUD_PROJECT" + value: "bigframes-load-testing" +} + +env_vars: { + key: "BIGFRAMES_TEST_MODEL_VERTEX_ENDPOINT" + value: "https://us-central1-aiplatform.googleapis.com/v1/projects/272725758477/locations/us-central1/endpoints/590545496255234048" +} diff --git a/.kokoro/load/notebook.cfg b/.kokoro/load/notebook.cfg new file mode 100644 index 00000000000..c14297019a1 --- /dev/null +++ b/.kokoro/load/notebook.cfg @@ -0,0 +1,17 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +# Only run this nox session. +env_vars: { + key: "NOX_SESSION" + value: "notebook" +} + +env_vars: { + key: "BENCHMARK_AND_PUBLISH" + value: "true" +} + +env_vars: { + key: "GOOGLE_CLOUD_PROJECT" + value: "bigframes-testing" +} diff --git a/.kokoro/populate-secrets.sh b/.kokoro/populate-secrets.sh new file mode 100755 index 00000000000..c435402f473 --- /dev/null +++ b/.kokoro/populate-secrets.sh @@ -0,0 +1,43 @@ +#!/bin/bash +# Copyright 2024 Google LLC. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +set -eo pipefail + +function now { date +"%Y-%m-%d %H:%M:%S" | tr -d '\n' ;} +function msg { println "$*" >&2 ;} +function println { printf '%s\n' "$(now) $*" ;} + + +# Populates requested secrets set in SECRET_MANAGER_KEYS from service account: +# kokoro-trampoline@cloud-devrel-kokoro-resources.iam.gserviceaccount.com +SECRET_LOCATION="${KOKORO_GFILE_DIR}/secret_manager" +msg "Creating folder on disk for secrets: ${SECRET_LOCATION}" +mkdir -p ${SECRET_LOCATION} +for key in $(echo ${SECRET_MANAGER_KEYS} | sed "s/,/ /g") +do + msg "Retrieving secret ${key}" + docker run --entrypoint=gcloud \ + --volume=${KOKORO_GFILE_DIR}:${KOKORO_GFILE_DIR} \ + gcr.io/google.com/cloudsdktool/cloud-sdk \ + secrets versions access latest \ + --project cloud-devrel-kokoro-resources \ + --secret ${key} > \ + "${SECRET_LOCATION}/${key}" + if [[ $? == 0 ]]; then + msg "Secret written to ${SECRET_LOCATION}/${key}" + else + msg "Error retrieving secret ${key}" + fi +done diff --git a/.kokoro/presubmit/common.cfg b/.kokoro/presubmit/common.cfg new file mode 100644 index 00000000000..97e0651aa92 --- /dev/null +++ b/.kokoro/presubmit/common.cfg @@ -0,0 +1,10 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +# Build logs will be here +action { + define_artifacts { + regex: "**/*sponge_log.xml" + } +} + +build_file: "python-bigquery-dataframes/.kokoro/build.sh" diff --git a/.kokoro/presubmit/doctest.cfg b/.kokoro/presubmit/doctest.cfg new file mode 100644 index 00000000000..2aad95beed4 --- /dev/null +++ b/.kokoro/presubmit/doctest.cfg @@ -0,0 +1,12 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +# Only run this nox session. +env_vars: { + key: "NOX_SESSION" + value: "cleanup doctest" +} + +env_vars: { + key: "GOOGLE_CLOUD_PROJECT" + value: "bigframes-testing" +} diff --git a/.kokoro/presubmit/e2e-gerrit.cfg b/.kokoro/presubmit/e2e-gerrit.cfg new file mode 100644 index 00000000000..19913344b6c --- /dev/null +++ b/.kokoro/presubmit/e2e-gerrit.cfg @@ -0,0 +1,7 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +# Only run this nox session. +env_vars: { + key: "NOX_SESSION" + value: "system_noextras e2e notebook" +} diff --git a/.kokoro/presubmit/e2e.cfg b/.kokoro/presubmit/e2e.cfg new file mode 100644 index 00000000000..e049dd30b3a --- /dev/null +++ b/.kokoro/presubmit/e2e.cfg @@ -0,0 +1,17 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +# Only run this nox session. +env_vars: { + key: "NOX_SESSION" + value: "e2e unit_prerelease system_prerelease system_noextras" +} + +env_vars: { + key: "GOOGLE_CLOUD_PROJECT" + value: "bigframes-load-testing" +} + +env_vars: { + key: "BIGFRAMES_TEST_MODEL_VERTEX_ENDPOINT" + value: "https://us-central1-aiplatform.googleapis.com/v1/projects/272725758477/locations/us-central1/endpoints/590545496255234048" +} diff --git a/.kokoro/presubmit/notebook.cfg b/.kokoro/presubmit/notebook.cfg new file mode 100644 index 00000000000..cc73c3bea40 --- /dev/null +++ b/.kokoro/presubmit/notebook.cfg @@ -0,0 +1,12 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +# Only run this nox session. +env_vars: { + key: "NOX_SESSION" + value: "notebook" +} + +env_vars: { + key: "GOOGLE_CLOUD_PROJECT" + value: "bigframes-testing" +} diff --git a/.kokoro/presubmit/prerelease-deps.cfg b/.kokoro/presubmit/prerelease-deps.cfg new file mode 100644 index 00000000000..3595fb43f5c --- /dev/null +++ b/.kokoro/presubmit/prerelease-deps.cfg @@ -0,0 +1,7 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +# Only run this nox session. +env_vars: { + key: "NOX_SESSION" + value: "prerelease_deps" +} diff --git a/.kokoro/presubmit/presubmit-gerrit.cfg b/.kokoro/presubmit/presubmit-gerrit.cfg new file mode 100644 index 00000000000..18a4c35325b --- /dev/null +++ b/.kokoro/presubmit/presubmit-gerrit.cfg @@ -0,0 +1 @@ +# Format: //devtools/kokoro/config/proto/build.proto diff --git a/.kokoro/presubmit/presubmit.cfg b/.kokoro/presubmit/presubmit.cfg new file mode 100644 index 00000000000..8f43917d92f --- /dev/null +++ b/.kokoro/presubmit/presubmit.cfg @@ -0,0 +1 @@ +# Format: //devtools/kokoro/config/proto/build.proto \ No newline at end of file diff --git a/.kokoro/presubmit/windows.cfg b/.kokoro/presubmit/windows.cfg new file mode 100644 index 00000000000..806986138dd --- /dev/null +++ b/.kokoro/presubmit/windows.cfg @@ -0,0 +1,3 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +build_file: "python-bigquery-dataframes/scripts/windows/build.bat" diff --git a/.kokoro/release-nightly.sh b/.kokoro/release-nightly.sh new file mode 100755 index 00000000000..df583f010d8 --- /dev/null +++ b/.kokoro/release-nightly.sh @@ -0,0 +1,124 @@ +#!/bin/bash +# Copyright 2020 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Based loosely on +# https://github.com/googleapis/python-bigquery/blob/main/.kokoro/release.sh + +set -eo pipefail +set -x + +# Parse command line arguments +DRY_RUN= +while [ $# -gt 0 ] ; do + case "$1" in + -d | --dry-run ) + DRY_RUN=true + ;; + -h | --help ) + echo -e "USAGE: `basename $0` [ -d | --dry-run ]" + exit + ;; + esac + shift 1; +done + +if [[ -z "${KOKORO_GOB_COMMIT}" ]]; then + PROJECT_SCM="github/python-bigquery-dataframes" +else + PROJECT_SCM="git/bigframes" +fi + +if [ -z "${PROJECT_ROOT:-}" ]; then + PROJECT_ROOT="${KOKORO_ARTIFACTS_DIR}/${PROJECT_SCM}" +fi + +# Move into the package, build the distribution and upload to shared bucket. +# See internal bug 274624240 for details. + +cd "${PROJECT_ROOT}" +rm -rf build dist + +# Workaround the fact that the repository that has been fetched before the +# build script. See: go/kokoro-native-docker-migration#known-issues and +# internal issue b/261050975. +git config --global --add safe.directory "${PROJECT_ROOT}" + +# Workaround for older pip not able to resolve dependencies. See internal +# issue 316909553. +python3.10 -m pip install pip==25.0.1 + +# Disable buffering, so that the logs stream through. +export PYTHONUNBUFFERED=1 + +# Install dependencies, as the following steps depend on it +python3.10 -m pip install -e .[all] + +# Update version string to include git hash and date +CURRENT_DATE=$(date '+%Y%m%d') +GIT_HASH=$(git rev-parse --short HEAD) +BIGFRAMES_VERSION=$(python3.10 -c "import bigframes; print(bigframes.__version__)") +RELEASE_VERSION=${BIGFRAMES_VERSION}dev${CURRENT_DATE}+${GIT_HASH} +sed -i -e "s/$BIGFRAMES_VERSION/$RELEASE_VERSION/g" bigframes/version.py + +# Generate the package wheel +python3.10 setup.py sdist bdist_wheel + +# Make sure that the wheel file is generated +VERSION_WHEEL=`ls dist/bigframes-*.whl` +num_wheel_files=`echo $VERSION_WHEEL | wc -w` +if [ $num_wheel_files -ne 1 ] ; then + echo "Exactly one wheel file should have been generated, found $num_wheel_files: $VERSION_WHEEL" + exit -1 +fi + +# Create a copy of the wheel with a well known, version agnostic name +LATEST_WHEEL=dist/bigframes-latest-py2.py3-none-any.whl +cp $VERSION_WHEEL $LATEST_WHEEL +cp dist/bigframes-*.tar.gz dist/bigframes-latest.tar.gz + +if ! [ ${DRY_RUN} ]; then +for gcs_path in gs://vertex_sdk_private_releases/bigframe/ \ + gs://dl-platform-colab/bigframes/ \ + gs://bigframes-wheels/; + do + gcloud storage cp --print-created-message dist/* ${gcs_path} + gcloud storage cp --print-created-message LICENSE ${gcs_path} + gcloud storage cp --recursive --print-created-message "notebooks/" ${gcs_path}notebooks/ + + done + + # publish API coverage information to BigQuery + # Note: only the kokoro service account has permission to write to this + # table, if you want to test this step, point it to a table you have + # write access to + COVERAGE_TABLE=bigframes-metrics.coverage_report.bigframes_coverage_nightly + python3.10 scripts/publish_api_coverage.py \ + bigquery \ + --bigframes_version=$BIGFRAMES_VERSION \ + --release_version=$RELEASE_VERSION \ + --bigquery_table=$COVERAGE_TABLE +fi + +# Undo the file changes, in case this script is running on a +# non-temporary instance of the bigframes repo +# TODO: This doesn't work with (set -eo pipefail) if the failure happened after +# the changes were made but before this cleanup, because the script would +# terminate with the failure itself. See if we can ensure the cleanup. +sed -i -e "s/$RELEASE_VERSION/$BIGFRAMES_VERSION/g" bigframes/version.py + +if ! [ ${DRY_RUN} ]; then + # Copy docs and wheels to Google Drive + python3.10 scripts/upload_to_google_drive.py +fi diff --git a/.kokoro/samples/lint/common.cfg b/.kokoro/samples/lint/common.cfg new file mode 100644 index 00000000000..b4d26c1f982 --- /dev/null +++ b/.kokoro/samples/lint/common.cfg @@ -0,0 +1,34 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +# Build logs will be here +action { + define_artifacts { + regex: "**/*sponge_log.xml" + } +} + +# Specify which tests to run +env_vars: { + key: "RUN_TESTS_SESSION" + value: "lint" +} + +env_vars: { + key: "TRAMPOLINE_BUILD_FILE" + value: "github/python-bigquery-dataframes/.kokoro/test-samples.sh" +} + +# Configure the docker image for kokoro-trampoline. +env_vars: { + key: "TRAMPOLINE_IMAGE" + value: "gcr.io/cloud-devrel-kokoro-resources/python-samples-testing-docker" +} + +# Download secrets for samples +gfile_resources: "/bigstore/cloud-devrel-kokoro-resources/python-docs-samples" + +# Download trampoline resources. +gfile_resources: "/bigstore/cloud-devrel-kokoro-resources/trampoline" + +# Use the trampoline script to run in docker. +build_file: "python-bigquery-dataframes/.kokoro/trampoline_v2.sh" \ No newline at end of file diff --git a/.kokoro/samples/lint/continuous.cfg b/.kokoro/samples/lint/continuous.cfg new file mode 100644 index 00000000000..a1c8d9759c8 --- /dev/null +++ b/.kokoro/samples/lint/continuous.cfg @@ -0,0 +1,6 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +env_vars: { + key: "INSTALL_LIBRARY_FROM_SOURCE" + value: "True" +} \ No newline at end of file diff --git a/.kokoro/samples/lint/periodic.cfg b/.kokoro/samples/lint/periodic.cfg new file mode 100644 index 00000000000..50fec964973 --- /dev/null +++ b/.kokoro/samples/lint/periodic.cfg @@ -0,0 +1,6 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +env_vars: { + key: "INSTALL_LIBRARY_FROM_SOURCE" + value: "False" +} \ No newline at end of file diff --git a/.kokoro/samples/lint/presubmit.cfg b/.kokoro/samples/lint/presubmit.cfg new file mode 100644 index 00000000000..a1c8d9759c8 --- /dev/null +++ b/.kokoro/samples/lint/presubmit.cfg @@ -0,0 +1,6 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +env_vars: { + key: "INSTALL_LIBRARY_FROM_SOURCE" + value: "True" +} \ No newline at end of file diff --git a/.kokoro/samples/python3.10/common.cfg b/.kokoro/samples/python3.10/common.cfg new file mode 100644 index 00000000000..8f9c66c571b --- /dev/null +++ b/.kokoro/samples/python3.10/common.cfg @@ -0,0 +1,40 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +# Build logs will be here +action { + define_artifacts { + regex: "**/*sponge_log.xml" + } +} + +# Specify which tests to run +env_vars: { + key: "RUN_TESTS_SESSION" + value: "py-3.10" +} + +# Declare build specific Cloud project. +env_vars: { + key: "BUILD_SPECIFIC_GCLOUD_PROJECT" + value: "python-docs-samples-tests-310" +} + +env_vars: { + key: "TRAMPOLINE_BUILD_FILE" + value: "github/python-bigquery-dataframes/.kokoro/test-samples.sh" +} + +# Configure the docker image for kokoro-trampoline. +env_vars: { + key: "TRAMPOLINE_IMAGE" + value: "gcr.io/cloud-devrel-kokoro-resources/python-samples-testing-docker" +} + +# Download secrets for samples +gfile_resources: "/bigstore/cloud-devrel-kokoro-resources/python-docs-samples" + +# Download trampoline resources. +gfile_resources: "/bigstore/cloud-devrel-kokoro-resources/trampoline" + +# Use the trampoline script to run in docker. +build_file: "python-bigquery-dataframes/.kokoro/trampoline_v2.sh" \ No newline at end of file diff --git a/.kokoro/samples/python3.10/continuous.cfg b/.kokoro/samples/python3.10/continuous.cfg new file mode 100644 index 00000000000..a1c8d9759c8 --- /dev/null +++ b/.kokoro/samples/python3.10/continuous.cfg @@ -0,0 +1,6 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +env_vars: { + key: "INSTALL_LIBRARY_FROM_SOURCE" + value: "True" +} \ No newline at end of file diff --git a/.kokoro/samples/python3.10/periodic-head.cfg b/.kokoro/samples/python3.10/periodic-head.cfg new file mode 100644 index 00000000000..123a35fbd3d --- /dev/null +++ b/.kokoro/samples/python3.10/periodic-head.cfg @@ -0,0 +1,11 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +env_vars: { + key: "INSTALL_LIBRARY_FROM_SOURCE" + value: "True" +} + +env_vars: { + key: "TRAMPOLINE_BUILD_FILE" + value: "github/python-bigquery-dataframes/.kokoro/test-samples-against-head.sh" +} diff --git a/.kokoro/samples/python3.10/periodic.cfg b/.kokoro/samples/python3.10/periodic.cfg new file mode 100644 index 00000000000..71cd1e597e3 --- /dev/null +++ b/.kokoro/samples/python3.10/periodic.cfg @@ -0,0 +1,6 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +env_vars: { + key: "INSTALL_LIBRARY_FROM_SOURCE" + value: "False" +} diff --git a/.kokoro/samples/python3.10/presubmit.cfg b/.kokoro/samples/python3.10/presubmit.cfg new file mode 100644 index 00000000000..a1c8d9759c8 --- /dev/null +++ b/.kokoro/samples/python3.10/presubmit.cfg @@ -0,0 +1,6 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +env_vars: { + key: "INSTALL_LIBRARY_FROM_SOURCE" + value: "True" +} \ No newline at end of file diff --git a/.kokoro/samples/python3.11/common.cfg b/.kokoro/samples/python3.11/common.cfg new file mode 100644 index 00000000000..1bba39114aa --- /dev/null +++ b/.kokoro/samples/python3.11/common.cfg @@ -0,0 +1,40 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +# Build logs will be here +action { + define_artifacts { + regex: "**/*sponge_log.xml" + } +} + +# Specify which tests to run +env_vars: { + key: "RUN_TESTS_SESSION" + value: "py-3.11" +} + +# Declare build specific Cloud project. +env_vars: { + key: "BUILD_SPECIFIC_GCLOUD_PROJECT" + value: "python-docs-samples-tests-311" +} + +env_vars: { + key: "TRAMPOLINE_BUILD_FILE" + value: "github/python-bigquery-dataframes/.kokoro/test-samples.sh" +} + +# Configure the docker image for kokoro-trampoline. +env_vars: { + key: "TRAMPOLINE_IMAGE" + value: "gcr.io/cloud-devrel-kokoro-resources/python-samples-testing-docker" +} + +# Download secrets for samples +gfile_resources: "/bigstore/cloud-devrel-kokoro-resources/python-docs-samples" + +# Download trampoline resources. +gfile_resources: "/bigstore/cloud-devrel-kokoro-resources/trampoline" + +# Use the trampoline script to run in docker. +build_file: "python-bigquery-dataframes/.kokoro/trampoline_v2.sh" \ No newline at end of file diff --git a/.kokoro/samples/python3.11/continuous.cfg b/.kokoro/samples/python3.11/continuous.cfg new file mode 100644 index 00000000000..a1c8d9759c8 --- /dev/null +++ b/.kokoro/samples/python3.11/continuous.cfg @@ -0,0 +1,6 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +env_vars: { + key: "INSTALL_LIBRARY_FROM_SOURCE" + value: "True" +} \ No newline at end of file diff --git a/.kokoro/samples/python3.11/periodic-head.cfg b/.kokoro/samples/python3.11/periodic-head.cfg new file mode 100644 index 00000000000..123a35fbd3d --- /dev/null +++ b/.kokoro/samples/python3.11/periodic-head.cfg @@ -0,0 +1,11 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +env_vars: { + key: "INSTALL_LIBRARY_FROM_SOURCE" + value: "True" +} + +env_vars: { + key: "TRAMPOLINE_BUILD_FILE" + value: "github/python-bigquery-dataframes/.kokoro/test-samples-against-head.sh" +} diff --git a/.kokoro/samples/python3.11/periodic.cfg b/.kokoro/samples/python3.11/periodic.cfg new file mode 100644 index 00000000000..71cd1e597e3 --- /dev/null +++ b/.kokoro/samples/python3.11/periodic.cfg @@ -0,0 +1,6 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +env_vars: { + key: "INSTALL_LIBRARY_FROM_SOURCE" + value: "False" +} diff --git a/.kokoro/samples/python3.11/presubmit.cfg b/.kokoro/samples/python3.11/presubmit.cfg new file mode 100644 index 00000000000..a1c8d9759c8 --- /dev/null +++ b/.kokoro/samples/python3.11/presubmit.cfg @@ -0,0 +1,6 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +env_vars: { + key: "INSTALL_LIBRARY_FROM_SOURCE" + value: "True" +} \ No newline at end of file diff --git a/.kokoro/samples/python3.12/common.cfg b/.kokoro/samples/python3.12/common.cfg new file mode 100644 index 00000000000..abf83e196db --- /dev/null +++ b/.kokoro/samples/python3.12/common.cfg @@ -0,0 +1,40 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +# Build logs will be here +action { + define_artifacts { + regex: "**/*sponge_log.xml" + } +} + +# Specify which tests to run +env_vars: { + key: "RUN_TESTS_SESSION" + value: "py-3.12" +} + +# Declare build specific Cloud project. +env_vars: { + key: "BUILD_SPECIFIC_GCLOUD_PROJECT" + value: "python-docs-samples-tests-312" +} + +env_vars: { + key: "TRAMPOLINE_BUILD_FILE" + value: "github/python-bigquery-dataframes/.kokoro/test-samples.sh" +} + +# Configure the docker image for kokoro-trampoline. +env_vars: { + key: "TRAMPOLINE_IMAGE" + value: "gcr.io/cloud-devrel-kokoro-resources/python-samples-testing-docker" +} + +# Download secrets for samples +gfile_resources: "/bigstore/cloud-devrel-kokoro-resources/python-docs-samples" + +# Download trampoline resources. +gfile_resources: "/bigstore/cloud-devrel-kokoro-resources/trampoline" + +# Use the trampoline script to run in docker. +build_file: "python-bigquery-dataframes/.kokoro/trampoline_v2.sh" \ No newline at end of file diff --git a/.kokoro/samples/python3.12/continuous.cfg b/.kokoro/samples/python3.12/continuous.cfg new file mode 100644 index 00000000000..a1c8d9759c8 --- /dev/null +++ b/.kokoro/samples/python3.12/continuous.cfg @@ -0,0 +1,6 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +env_vars: { + key: "INSTALL_LIBRARY_FROM_SOURCE" + value: "True" +} \ No newline at end of file diff --git a/.kokoro/samples/python3.12/periodic-head.cfg b/.kokoro/samples/python3.12/periodic-head.cfg new file mode 100644 index 00000000000..123a35fbd3d --- /dev/null +++ b/.kokoro/samples/python3.12/periodic-head.cfg @@ -0,0 +1,11 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +env_vars: { + key: "INSTALL_LIBRARY_FROM_SOURCE" + value: "True" +} + +env_vars: { + key: "TRAMPOLINE_BUILD_FILE" + value: "github/python-bigquery-dataframes/.kokoro/test-samples-against-head.sh" +} diff --git a/.kokoro/samples/python3.12/periodic.cfg b/.kokoro/samples/python3.12/periodic.cfg new file mode 100644 index 00000000000..71cd1e597e3 --- /dev/null +++ b/.kokoro/samples/python3.12/periodic.cfg @@ -0,0 +1,6 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +env_vars: { + key: "INSTALL_LIBRARY_FROM_SOURCE" + value: "False" +} diff --git a/.kokoro/samples/python3.12/presubmit.cfg b/.kokoro/samples/python3.12/presubmit.cfg new file mode 100644 index 00000000000..a1c8d9759c8 --- /dev/null +++ b/.kokoro/samples/python3.12/presubmit.cfg @@ -0,0 +1,6 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +env_vars: { + key: "INSTALL_LIBRARY_FROM_SOURCE" + value: "True" +} \ No newline at end of file diff --git a/.kokoro/samples/python3.13/common.cfg b/.kokoro/samples/python3.13/common.cfg new file mode 100644 index 00000000000..6a5d9a20800 --- /dev/null +++ b/.kokoro/samples/python3.13/common.cfg @@ -0,0 +1,40 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +# Build logs will be here +action { + define_artifacts { + regex: "**/*sponge_log.xml" + } +} + +# Specify which tests to run +env_vars: { + key: "RUN_TESTS_SESSION" + value: "py-3.13" +} + +# Declare build specific Cloud project. +env_vars: { + key: "BUILD_SPECIFIC_GCLOUD_PROJECT" + value: "python-docs-samples-tests-313" +} + +env_vars: { + key: "TRAMPOLINE_BUILD_FILE" + value: "github/python-bigquery-dataframes/.kokoro/test-samples.sh" +} + +# Configure the docker image for kokoro-trampoline. +env_vars: { + key: "TRAMPOLINE_IMAGE" + value: "gcr.io/cloud-devrel-kokoro-resources/python-samples-testing-docker" +} + +# Download secrets for samples +gfile_resources: "/bigstore/cloud-devrel-kokoro-resources/python-docs-samples" + +# Download trampoline resources. +gfile_resources: "/bigstore/cloud-devrel-kokoro-resources/trampoline" + +# Use the trampoline script to run in docker. +build_file: "python-bigquery-dataframes/.kokoro/trampoline_v2.sh" diff --git a/.kokoro/samples/python3.13/continuous.cfg b/.kokoro/samples/python3.13/continuous.cfg new file mode 100644 index 00000000000..a1c8d9759c8 --- /dev/null +++ b/.kokoro/samples/python3.13/continuous.cfg @@ -0,0 +1,6 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +env_vars: { + key: "INSTALL_LIBRARY_FROM_SOURCE" + value: "True" +} \ No newline at end of file diff --git a/.kokoro/samples/python3.13/periodic-head.cfg b/.kokoro/samples/python3.13/periodic-head.cfg new file mode 100644 index 00000000000..123a35fbd3d --- /dev/null +++ b/.kokoro/samples/python3.13/periodic-head.cfg @@ -0,0 +1,11 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +env_vars: { + key: "INSTALL_LIBRARY_FROM_SOURCE" + value: "True" +} + +env_vars: { + key: "TRAMPOLINE_BUILD_FILE" + value: "github/python-bigquery-dataframes/.kokoro/test-samples-against-head.sh" +} diff --git a/.kokoro/samples/python3.13/periodic.cfg b/.kokoro/samples/python3.13/periodic.cfg new file mode 100644 index 00000000000..71cd1e597e3 --- /dev/null +++ b/.kokoro/samples/python3.13/periodic.cfg @@ -0,0 +1,6 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +env_vars: { + key: "INSTALL_LIBRARY_FROM_SOURCE" + value: "False" +} diff --git a/.kokoro/samples/python3.13/presubmit.cfg b/.kokoro/samples/python3.13/presubmit.cfg new file mode 100644 index 00000000000..a1c8d9759c8 --- /dev/null +++ b/.kokoro/samples/python3.13/presubmit.cfg @@ -0,0 +1,6 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +env_vars: { + key: "INSTALL_LIBRARY_FROM_SOURCE" + value: "True" +} \ No newline at end of file diff --git a/.kokoro/test-samples-against-head.sh b/.kokoro/test-samples-against-head.sh new file mode 100755 index 00000000000..e9d8bd79a64 --- /dev/null +++ b/.kokoro/test-samples-against-head.sh @@ -0,0 +1,26 @@ +#!/bin/bash +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# A customized test runner for samples. +# +# For periodic builds, you can specify this file for testing against head. + +# `-e` enables the script to automatically fail when a command fails +# `-o pipefail` sets the exit code to the rightmost comment to exit with a non-zero +set -eo pipefail +# Enables `**` to include files nested inside sub-folders +shopt -s globstar + +exec .kokoro/test-samples-impl.sh diff --git a/.kokoro/test-samples-impl.sh b/.kokoro/test-samples-impl.sh new file mode 100755 index 00000000000..97cdc9c13fe --- /dev/null +++ b/.kokoro/test-samples-impl.sh @@ -0,0 +1,103 @@ +#!/bin/bash +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +# `-e` enables the script to automatically fail when a command fails +# `-o pipefail` sets the exit code to the rightmost comment to exit with a non-zero +set -eo pipefail +# Enables `**` to include files nested inside sub-folders +shopt -s globstar + +# Exit early if samples don't exist +if ! find samples -name 'requirements.txt' | grep -q .; then + echo "No tests run. './samples/**/requirements.txt' not found" + exit 0 +fi + +# Disable buffering, so that the logs stream through. +export PYTHONUNBUFFERED=1 + +# Debug: show build environment +env | grep KOKORO + +# Install nox +# `virtualenv==20.26.6` is added for Python 3.7 compatibility +python3.10 -m pip install --upgrade --quiet nox virtualenv==20.26.6 + +# Use secrets acessor service account to get secrets +if [[ -f "${KOKORO_GFILE_DIR}/secrets_viewer_service_account.json" ]]; then + gcloud auth activate-service-account \ + --key-file="${KOKORO_GFILE_DIR}/secrets_viewer_service_account.json" \ + --project="cloud-devrel-kokoro-resources" +fi + +# This script will create 3 files: +# - testing/test-env.sh +# - testing/service-account.json +# - testing/client-secrets.json +./scripts/decrypt-secrets.sh + +source ./testing/test-env.sh +export GOOGLE_APPLICATION_CREDENTIALS=$(pwd)/testing/service-account.json + +# For cloud-run session, we activate the service account for gcloud sdk. +gcloud auth activate-service-account \ + --key-file "${GOOGLE_APPLICATION_CREDENTIALS}" + +export GOOGLE_CLIENT_SECRETS=$(pwd)/testing/client-secrets.json + +echo -e "\n******************** TESTING PROJECTS ********************" + +# Switch to 'fail at end' to allow all tests to complete before exiting. +set +e +# Use RTN to return a non-zero value if the test fails. +RTN=0 +ROOT=$(pwd) +# Find all requirements.txt in the samples directory (may break on whitespace). +for file in samples/**/requirements.txt; do + cd "$ROOT" + # Navigate to the project folder. + file=$(dirname "$file") + cd "$file" + + echo "------------------------------------------------------------" + echo "- testing $file" + echo "------------------------------------------------------------" + + # Use nox to execute the tests for the project. + python3.10 -m nox -s "$RUN_TESTS_SESSION" + EXIT=$? + + # If this is a periodic build, send the test log to the FlakyBot. + # See https://github.com/googleapis/repo-automation-bots/tree/main/packages/flakybot. + if [[ $KOKORO_BUILD_ARTIFACTS_SUBDIR = *"periodic"* ]]; then + chmod +x $KOKORO_GFILE_DIR/linux_amd64/flakybot + $KOKORO_GFILE_DIR/linux_amd64/flakybot + fi + + if [[ $EXIT -ne 0 ]]; then + RTN=1 + echo -e "\n Testing failed: Nox returned a non-zero exit code. \n" + else + echo -e "\n Testing completed.\n" + fi + +done +cd "$ROOT" + +# Workaround for Kokoro permissions issue: delete secrets +rm testing/{test-env.sh,client-secrets.json,service-account.json} + +exit "$RTN" diff --git a/.kokoro/test-samples.sh b/.kokoro/test-samples.sh new file mode 100755 index 00000000000..7933d820149 --- /dev/null +++ b/.kokoro/test-samples.sh @@ -0,0 +1,44 @@ +#!/bin/bash +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# The default test runner for samples. +# +# For periodic builds, we rewinds the repo to the latest release, and +# run test-samples-impl.sh. + +# `-e` enables the script to automatically fail when a command fails +# `-o pipefail` sets the exit code to the rightmost comment to exit with a non-zero +set -eo pipefail +# Enables `**` to include files nested inside sub-folders +shopt -s globstar + +# Run periodic samples tests at latest release +if [[ $KOKORO_BUILD_ARTIFACTS_SUBDIR = *"periodic"* ]]; then + # preserving the test runner implementation. + cp .kokoro/test-samples-impl.sh "${TMPDIR}/test-samples-impl.sh" + echo "--- IMPORTANT IMPORTANT IMPORTANT ---" + echo "Now we rewind the repo back to the latest release..." + LATEST_RELEASE=$(git describe --abbrev=0 --tags) + git checkout $LATEST_RELEASE + echo "The current head is: " + echo $(git rev-parse --verify HEAD) + echo "--- IMPORTANT IMPORTANT IMPORTANT ---" + # move back the test runner implementation if there's no file. + if [ ! -f .kokoro/test-samples-impl.sh ]; then + cp "${TMPDIR}/test-samples-impl.sh" .kokoro/test-samples-impl.sh + fi +fi + +exec .kokoro/test-samples-impl.sh diff --git a/.kokoro/trampoline.sh b/.kokoro/trampoline.sh new file mode 100755 index 00000000000..48f79699706 --- /dev/null +++ b/.kokoro/trampoline.sh @@ -0,0 +1,28 @@ +#!/bin/bash +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +set -eo pipefail + +# Always run the cleanup script, regardless of the success of bouncing into +# the container. +function cleanup() { + chmod +x ${KOKORO_GFILE_DIR}/trampoline_cleanup.sh + ${KOKORO_GFILE_DIR}/trampoline_cleanup.sh + echo "cleanup"; +} +trap cleanup EXIT + +$(dirname $0)/populate-secrets.sh # Secret Manager secrets. +python3 "${KOKORO_GFILE_DIR}/trampoline_v1.py" \ No newline at end of file diff --git a/.kokoro/trampoline_v2.sh b/.kokoro/trampoline_v2.sh new file mode 100755 index 00000000000..d03f92dfc48 --- /dev/null +++ b/.kokoro/trampoline_v2.sh @@ -0,0 +1,487 @@ +#!/usr/bin/env bash +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# trampoline_v2.sh +# +# This script does 3 things. +# +# 1. Prepare the Docker image for the test +# 2. Run the Docker with appropriate flags to run the test +# 3. Upload the newly built Docker image +# +# in a way that is somewhat compatible with trampoline_v1. +# +# To run this script, first download few files from gcs to /dev/shm. +# (/dev/shm is passed into the container as KOKORO_GFILE_DIR). +# +# gcloud storage cp gs://cloud-devrel-kokoro-resources/python-docs-samples/secrets_viewer_service_account.json /dev/shm +# gcloud storage cp gs://cloud-devrel-kokoro-resources/python-docs-samples/automl_secrets.txt /dev/shm +# +# Then run the script. +# .kokoro/trampoline_v2.sh +# +# These environment variables are required: +# TRAMPOLINE_IMAGE: The docker image to use. +# TRAMPOLINE_DOCKERFILE: The location of the Dockerfile. +# +# You can optionally change these environment variables: +# TRAMPOLINE_IMAGE_UPLOAD: +# (true|false): Whether to upload the Docker image after the +# successful builds. +# TRAMPOLINE_BUILD_FILE: The script to run in the docker container. +# TRAMPOLINE_WORKSPACE: The workspace path in the docker container. +# Defaults to /workspace. +# Potentially there are some repo specific envvars in .trampolinerc in +# the project root. + + +set -euo pipefail + +TRAMPOLINE_VERSION="2.0.5" + +if command -v tput >/dev/null && [[ -n "${TERM:-}" ]]; then + readonly IO_COLOR_RED="$(tput setaf 1)" + readonly IO_COLOR_GREEN="$(tput setaf 2)" + readonly IO_COLOR_YELLOW="$(tput setaf 3)" + readonly IO_COLOR_RESET="$(tput sgr0)" +else + readonly IO_COLOR_RED="" + readonly IO_COLOR_GREEN="" + readonly IO_COLOR_YELLOW="" + readonly IO_COLOR_RESET="" +fi + +function function_exists { + [ $(LC_ALL=C type -t $1)"" == "function" ] +} + +# Logs a message using the given color. The first argument must be one +# of the IO_COLOR_* variables defined above, such as +# "${IO_COLOR_YELLOW}". The remaining arguments will be logged in the +# given color. The log message will also have an RFC-3339 timestamp +# prepended (in UTC). You can disable the color output by setting +# TERM=vt100. +function log_impl() { + local color="$1" + shift + local timestamp="$(date -u "+%Y-%m-%dT%H:%M:%SZ")" + echo "================================================================" + echo "${color}${timestamp}:" "$@" "${IO_COLOR_RESET}" + echo "================================================================" +} + +# Logs the given message with normal coloring and a timestamp. +function log() { + log_impl "${IO_COLOR_RESET}" "$@" +} + +# Logs the given message in green with a timestamp. +function log_green() { + log_impl "${IO_COLOR_GREEN}" "$@" +} + +# Logs the given message in yellow with a timestamp. +function log_yellow() { + log_impl "${IO_COLOR_YELLOW}" "$@" +} + +# Logs the given message in red with a timestamp. +function log_red() { + log_impl "${IO_COLOR_RED}" "$@" +} + +readonly tmpdir=$(mktemp -d -t ci-XXXXXXXX) +readonly tmphome="${tmpdir}/h" +mkdir -p "${tmphome}" + +function cleanup() { + rm -rf "${tmpdir}" +} +trap cleanup EXIT + +RUNNING_IN_CI="${RUNNING_IN_CI:-false}" + +# The workspace in the container, defaults to /workspace. +TRAMPOLINE_WORKSPACE="${TRAMPOLINE_WORKSPACE:-/workspace}" + +pass_down_envvars=( + # TRAMPOLINE_V2 variables. + # Tells scripts whether they are running as part of CI or not. + "RUNNING_IN_CI" + # Indicates which CI system we're in. + "TRAMPOLINE_CI" + # Indicates the version of the script. + "TRAMPOLINE_VERSION" +) + +log_yellow "Building with Trampoline ${TRAMPOLINE_VERSION}" + +# Detect which CI systems we're in. If we're in any of the CI systems +# we support, `RUNNING_IN_CI` will be true and `TRAMPOLINE_CI` will be +# the name of the CI system. Both envvars will be passing down to the +# container for telling which CI system we're in. +if [[ -n "${KOKORO_BUILD_ID:-}" ]]; then + # descriptive env var for indicating it's on CI. + RUNNING_IN_CI="true" + TRAMPOLINE_CI="kokoro" + if [[ "${TRAMPOLINE_USE_LEGACY_SERVICE_ACCOUNT:-}" == "true" ]]; then + if [[ ! -f "${KOKORO_GFILE_DIR}/kokoro-trampoline.service-account.json" ]]; then + log_red "${KOKORO_GFILE_DIR}/kokoro-trampoline.service-account.json does not exist. Did you forget to mount cloud-devrel-kokoro-resources/trampoline? Aborting." + exit 1 + fi + # This service account will be activated later. + TRAMPOLINE_SERVICE_ACCOUNT="${KOKORO_GFILE_DIR}/kokoro-trampoline.service-account.json" + else + if [[ "${TRAMPOLINE_VERBOSE:-}" == "true" ]]; then + gcloud auth list + fi + log_yellow "Configuring Container Registry access" + gcloud auth configure-docker --quiet + fi + pass_down_envvars+=( + # KOKORO dynamic variables. + "KOKORO_BUILD_NUMBER" + "KOKORO_BUILD_ID" + "KOKORO_JOB_NAME" + "KOKORO_GIT_COMMIT" + "KOKORO_GITHUB_COMMIT" + "KOKORO_GITHUB_PULL_REQUEST_NUMBER" + "KOKORO_GITHUB_PULL_REQUEST_COMMIT" + # For FlakyBot + "KOKORO_GITHUB_COMMIT_URL" + "KOKORO_GITHUB_PULL_REQUEST_URL" + ) +elif [[ "${TRAVIS:-}" == "true" ]]; then + RUNNING_IN_CI="true" + TRAMPOLINE_CI="travis" + pass_down_envvars+=( + "TRAVIS_BRANCH" + "TRAVIS_BUILD_ID" + "TRAVIS_BUILD_NUMBER" + "TRAVIS_BUILD_WEB_URL" + "TRAVIS_COMMIT" + "TRAVIS_COMMIT_MESSAGE" + "TRAVIS_COMMIT_RANGE" + "TRAVIS_JOB_NAME" + "TRAVIS_JOB_NUMBER" + "TRAVIS_JOB_WEB_URL" + "TRAVIS_PULL_REQUEST" + "TRAVIS_PULL_REQUEST_BRANCH" + "TRAVIS_PULL_REQUEST_SHA" + "TRAVIS_PULL_REQUEST_SLUG" + "TRAVIS_REPO_SLUG" + "TRAVIS_SECURE_ENV_VARS" + "TRAVIS_TAG" + ) +elif [[ -n "${GITHUB_RUN_ID:-}" ]]; then + RUNNING_IN_CI="true" + TRAMPOLINE_CI="github-workflow" + pass_down_envvars+=( + "GITHUB_WORKFLOW" + "GITHUB_RUN_ID" + "GITHUB_RUN_NUMBER" + "GITHUB_ACTION" + "GITHUB_ACTIONS" + "GITHUB_ACTOR" + "GITHUB_REPOSITORY" + "GITHUB_EVENT_NAME" + "GITHUB_EVENT_PATH" + "GITHUB_SHA" + "GITHUB_REF" + "GITHUB_HEAD_REF" + "GITHUB_BASE_REF" + ) +elif [[ "${CIRCLECI:-}" == "true" ]]; then + RUNNING_IN_CI="true" + TRAMPOLINE_CI="circleci" + pass_down_envvars+=( + "CIRCLE_BRANCH" + "CIRCLE_BUILD_NUM" + "CIRCLE_BUILD_URL" + "CIRCLE_COMPARE_URL" + "CIRCLE_JOB" + "CIRCLE_NODE_INDEX" + "CIRCLE_NODE_TOTAL" + "CIRCLE_PREVIOUS_BUILD_NUM" + "CIRCLE_PROJECT_REPONAME" + "CIRCLE_PROJECT_USERNAME" + "CIRCLE_REPOSITORY_URL" + "CIRCLE_SHA1" + "CIRCLE_STAGE" + "CIRCLE_USERNAME" + "CIRCLE_WORKFLOW_ID" + "CIRCLE_WORKFLOW_JOB_ID" + "CIRCLE_WORKFLOW_UPSTREAM_JOB_IDS" + "CIRCLE_WORKFLOW_WORKSPACE_ID" + ) +fi + +# Configure the service account for pulling the docker image. +function repo_root() { + local dir="$1" + while [[ ! -d "${dir}/.git" ]]; do + dir="$(dirname "$dir")" + done + echo "${dir}" +} + +# Detect the project root. In CI builds, we assume the script is in +# the git tree and traverse from there, otherwise, traverse from `pwd` +# to find `.git` directory. +if [[ "${RUNNING_IN_CI:-}" == "true" ]]; then + PROGRAM_PATH="$(realpath "$0")" + PROGRAM_DIR="$(dirname "${PROGRAM_PATH}")" + PROJECT_ROOT="$(repo_root "${PROGRAM_DIR}")" +else + PROJECT_ROOT="$(repo_root $(pwd))" +fi + +log_yellow "Changing to the project root: ${PROJECT_ROOT}." +cd "${PROJECT_ROOT}" + +# To support relative path for `TRAMPOLINE_SERVICE_ACCOUNT`, we need +# to use this environment variable in `PROJECT_ROOT`. +if [[ -n "${TRAMPOLINE_SERVICE_ACCOUNT:-}" ]]; then + + mkdir -p "${tmpdir}/gcloud" + gcloud_config_dir="${tmpdir}/gcloud" + + log_yellow "Using isolated gcloud config: ${gcloud_config_dir}." + export CLOUDSDK_CONFIG="${gcloud_config_dir}" + + log_yellow "Using ${TRAMPOLINE_SERVICE_ACCOUNT} for authentication." + gcloud auth activate-service-account \ + --key-file "${TRAMPOLINE_SERVICE_ACCOUNT}" + log_yellow "Configuring Container Registry access" + gcloud auth configure-docker --quiet +fi + +required_envvars=( + # The basic trampoline configurations. + "TRAMPOLINE_IMAGE" + "TRAMPOLINE_BUILD_FILE" +) + +if [[ -f "${PROJECT_ROOT}/.trampolinerc" ]]; then + source "${PROJECT_ROOT}/.trampolinerc" +fi + +log_yellow "Checking environment variables." +for e in "${required_envvars[@]}" +do + if [[ -z "${!e:-}" ]]; then + log "Missing ${e} env var. Aborting." + exit 1 + fi +done + +# We want to support legacy style TRAMPOLINE_BUILD_FILE used with V1 +# script: e.g. "github/repo-name/.kokoro/run_tests.sh" +TRAMPOLINE_BUILD_FILE="${TRAMPOLINE_BUILD_FILE#github/*/}" +log_yellow "Using TRAMPOLINE_BUILD_FILE: ${TRAMPOLINE_BUILD_FILE}" + +# ignore error on docker operations and test execution +set +e + +log_yellow "Preparing Docker image." +# We only download the docker image in CI builds. +if [[ "${RUNNING_IN_CI:-}" == "true" ]]; then + # Download the docker image specified by `TRAMPOLINE_IMAGE` + + # We may want to add --max-concurrent-downloads flag. + + log_yellow "Start pulling the Docker image: ${TRAMPOLINE_IMAGE}." + if docker pull "${TRAMPOLINE_IMAGE}"; then + log_green "Finished pulling the Docker image: ${TRAMPOLINE_IMAGE}." + has_image="true" + else + log_red "Failed pulling the Docker image: ${TRAMPOLINE_IMAGE}." + has_image="false" + fi +else + # For local run, check if we have the image. + if docker images "${TRAMPOLINE_IMAGE}:latest" | grep "${TRAMPOLINE_IMAGE}"; then + has_image="true" + else + has_image="false" + fi +fi + + +# The default user for a Docker container has uid 0 (root). To avoid +# creating root-owned files in the build directory we tell docker to +# use the current user ID. +user_uid="$(id -u)" +user_gid="$(id -g)" +user_name="$(id -un)" + +# To allow docker in docker, we add the user to the docker group in +# the host os. +docker_gid=$(cut -d: -f3 < <(getent group docker)) + +update_cache="false" +if [[ "${TRAMPOLINE_DOCKERFILE:-none}" != "none" ]]; then + # Build the Docker image from the source. + context_dir=$(dirname "${TRAMPOLINE_DOCKERFILE}") + docker_build_flags=( + "-f" "${TRAMPOLINE_DOCKERFILE}" + "-t" "${TRAMPOLINE_IMAGE}" + "--build-arg" "UID=${user_uid}" + "--build-arg" "USERNAME=${user_name}" + ) + if [[ "${has_image}" == "true" ]]; then + docker_build_flags+=("--cache-from" "${TRAMPOLINE_IMAGE}") + fi + + log_yellow "Start building the docker image." + if [[ "${TRAMPOLINE_VERBOSE:-false}" == "true" ]]; then + echo "docker build" "${docker_build_flags[@]}" "${context_dir}" + fi + + # ON CI systems, we want to suppress docker build logs, only + # output the logs when it fails. + if [[ "${RUNNING_IN_CI:-}" == "true" ]]; then + if docker build "${docker_build_flags[@]}" "${context_dir}" \ + > "${tmpdir}/docker_build.log" 2>&1; then + if [[ "${TRAMPOLINE_VERBOSE:-}" == "true" ]]; then + cat "${tmpdir}/docker_build.log" + fi + + log_green "Finished building the docker image." + update_cache="true" + else + log_red "Failed to build the Docker image, aborting." + log_yellow "Dumping the build logs:" + cat "${tmpdir}/docker_build.log" + exit 1 + fi + else + if docker build "${docker_build_flags[@]}" "${context_dir}"; then + log_green "Finished building the docker image." + update_cache="true" + else + log_red "Failed to build the Docker image, aborting." + exit 1 + fi + fi +else + if [[ "${has_image}" != "true" ]]; then + log_red "We do not have ${TRAMPOLINE_IMAGE} locally, aborting." + exit 1 + fi +fi + +# We use an array for the flags so they are easier to document. +docker_flags=( + # Remove the container after it exists. + "--rm" + + # Use the host network. + "--network=host" + + # Run in priviledged mode. We are not using docker for sandboxing or + # isolation, just for packaging our dev tools. + "--privileged" + + # Run the docker script with the user id. Because the docker image gets to + # write in ${PWD} you typically want this to be your user id. + # To allow docker in docker, we need to use docker gid on the host. + "--user" "${user_uid}:${docker_gid}" + + # Pass down the USER. + "--env" "USER=${user_name}" + + # Mount the project directory inside the Docker container. + "--volume" "${PROJECT_ROOT}:${TRAMPOLINE_WORKSPACE}" + "--workdir" "${TRAMPOLINE_WORKSPACE}" + "--env" "PROJECT_ROOT=${TRAMPOLINE_WORKSPACE}" + + # Mount the temporary home directory. + "--volume" "${tmphome}:/h" + "--env" "HOME=/h" + + # Allow docker in docker. + "--volume" "/var/run/docker.sock:/var/run/docker.sock" + + # Mount the /tmp so that docker in docker can mount the files + # there correctly. + "--volume" "/tmp:/tmp" + # Pass down the KOKORO_GFILE_DIR and KOKORO_KEYSTORE_DIR + # TODO(tmatsuo): This part is not portable. + "--env" "TRAMPOLINE_SECRET_DIR=/secrets" + "--volume" "${KOKORO_GFILE_DIR:-/dev/shm}:/secrets/gfile" + "--env" "KOKORO_GFILE_DIR=/secrets/gfile" + "--volume" "${KOKORO_KEYSTORE_DIR:-/dev/shm}:/secrets/keystore" + "--env" "KOKORO_KEYSTORE_DIR=/secrets/keystore" +) + +# Add an option for nicer output if the build gets a tty. +if [[ -t 0 ]]; then + docker_flags+=("-it") +fi + +# Passing down env vars +for e in "${pass_down_envvars[@]}" +do + if [[ -n "${!e:-}" ]]; then + docker_flags+=("--env" "${e}=${!e}") + fi +done + +# If arguments are given, all arguments will become the commands run +# in the container, otherwise run TRAMPOLINE_BUILD_FILE. +if [[ $# -ge 1 ]]; then + log_yellow "Running the given commands '" "${@:1}" "' in the container." + readonly commands=("${@:1}") + if [[ "${TRAMPOLINE_VERBOSE:-}" == "true" ]]; then + echo docker run "${docker_flags[@]}" "${TRAMPOLINE_IMAGE}" "${commands[@]}" + fi + docker run "${docker_flags[@]}" "${TRAMPOLINE_IMAGE}" "${commands[@]}" +else + log_yellow "Running the tests in a Docker container." + docker_flags+=("--entrypoint=${TRAMPOLINE_BUILD_FILE}") + if [[ "${TRAMPOLINE_VERBOSE:-}" == "true" ]]; then + echo docker run "${docker_flags[@]}" "${TRAMPOLINE_IMAGE}" + fi + docker run "${docker_flags[@]}" "${TRAMPOLINE_IMAGE}" +fi + + +test_retval=$? + +if [[ ${test_retval} -eq 0 ]]; then + log_green "Build finished with ${test_retval}" +else + log_red "Build finished with ${test_retval}" +fi + +# Only upload it when the test passes. +if [[ "${update_cache}" == "true" ]] && \ + [[ $test_retval == 0 ]] && \ + [[ "${TRAMPOLINE_IMAGE_UPLOAD:-false}" == "true" ]]; then + log_yellow "Uploading the Docker image." + if docker push "${TRAMPOLINE_IMAGE}"; then + log_green "Finished uploading the Docker image." + else + log_red "Failed uploading the Docker image." + fi + # Call trampoline_after_upload_hook if it's defined. + if function_exists trampoline_after_upload_hook; then + trampoline_after_upload_hook + fi + +fi + +exit "${test_retval}" diff --git a/.librarian/state.yaml b/.librarian/state.yaml new file mode 100644 index 00000000000..0344252107e --- /dev/null +++ b/.librarian/state.yaml @@ -0,0 +1,11 @@ +image: us-central1-docker.pkg.dev/cloud-sdk-librarian-prod/images-prod/python-librarian-generator@sha256:160860d189ff1c2f7515638478823712fa5b243e27ccc33a2728669fa1e2ed0c +libraries: + - id: bigframes + version: 2.39.0 + last_generated_commit: "" + apis: [] + source_roots: + - . + preserve_regex: [] + remove_regex: [] + tag_format: v{version} diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 00000000000..2dc978f0032 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,59 @@ +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# See https://pre-commit.com for more information +# See https://pre-commit.com/hooks.html for more hooks +default_install_hook_types: +- pre-commit +- commit-msg + +repos: +- repo: https://github.com/pre-commit/pre-commit-hooks + rev: v4.0.1 + hooks: + - id: trailing-whitespace + - id: end-of-file-fixer + exclude: "^tests/unit/core/compile/sqlglot/.*snapshots" + - id: check-yaml +- repo: https://github.com/pycqa/isort + rev: 5.12.0 + hooks: + - id: isort + name: isort (python) +- repo: https://github.com/psf/black + rev: 23.7.0 + hooks: + - id: black +- repo: https://github.com/pycqa/flake8 + rev: 7.1.2 + hooks: + - id: flake8 +- repo: https://github.com/pre-commit/mirrors-mypy + rev: v1.15.0 + hooks: + - id: mypy + additional_dependencies: [types-requests, types-tabulate, types-PyYAML, pandas-stubs<=2.2.3.241126] + exclude: "^third_party" + args: ["--check-untyped-defs", "--explicit-package-bases", "--ignore-missing-imports"] +- repo: https://github.com/biomejs/pre-commit + rev: v2.2.4 + hooks: + - id: biome-check + files: '\.(js|css)$' +- repo: https://github.com/compilerla/conventional-pre-commit + rev: fdde5f0251edbfc554795afdd6df71826d6602f3 + hooks: + - id: conventional-pre-commit + stages: [commit-msg] + args: [] diff --git a/.trampolinerc b/.trampolinerc new file mode 100644 index 00000000000..0080152373d --- /dev/null +++ b/.trampolinerc @@ -0,0 +1,61 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Add required env vars here. +required_envvars+=( +) + +# Add env vars which are passed down into the container here. +pass_down_envvars+=( + "NOX_SESSION" + ############### + # Docs builds + ############### + "STAGING_BUCKET" + "V2_STAGING_BUCKET" + ################## + # Samples builds + ################## + "INSTALL_LIBRARY_FROM_SOURCE" + "RUN_TESTS_SESSION" + "BUILD_SPECIFIC_GCLOUD_PROJECT" + # Target directories. + "RUN_TESTS_DIRS" + # The nox session to run. + "RUN_TESTS_SESSION" +) + +# Prevent unintentional override on the default image. +if [[ "${TRAMPOLINE_IMAGE_UPLOAD:-false}" == "true" ]] && \ + [[ -z "${TRAMPOLINE_IMAGE:-}" ]]; then + echo "Please set TRAMPOLINE_IMAGE if you want to upload the Docker image." + exit 1 +fi + +# Define the default value if it makes sense. +if [[ -z "${TRAMPOLINE_IMAGE_UPLOAD:-}" ]]; then + TRAMPOLINE_IMAGE_UPLOAD="" +fi + +if [[ -z "${TRAMPOLINE_IMAGE:-}" ]]; then + TRAMPOLINE_IMAGE="" +fi + +if [[ -z "${TRAMPOLINE_DOCKERFILE:-}" ]]; then + TRAMPOLINE_DOCKERFILE="" +fi + +if [[ -z "${TRAMPOLINE_BUILD_FILE:-}" ]]; then + TRAMPOLINE_BUILD_FILE="" +fi diff --git a/CODE_OF_CONDUCT.md b/CODE_OF_CONDUCT.md new file mode 100644 index 00000000000..039f4368120 --- /dev/null +++ b/CODE_OF_CONDUCT.md @@ -0,0 +1,95 @@ + +# Code of Conduct + +## Our Pledge + +In the interest of fostering an open and welcoming environment, we as +contributors and maintainers pledge to making participation in our project and +our community a harassment-free experience for everyone, regardless of age, body +size, disability, ethnicity, gender identity and expression, level of +experience, education, socio-economic status, nationality, personal appearance, +race, religion, or sexual identity and orientation. + +## Our Standards + +Examples of behavior that contributes to creating a positive environment +include: + +* Using welcoming and inclusive language +* Being respectful of differing viewpoints and experiences +* Gracefully accepting constructive criticism +* Focusing on what is best for the community +* Showing empathy towards other community members + +Examples of unacceptable behavior by participants include: + +* The use of sexualized language or imagery and unwelcome sexual attention or + advances +* Trolling, insulting/derogatory comments, and personal or political attacks +* Public or private harassment +* Publishing others' private information, such as a physical or electronic + address, without explicit permission +* Other conduct which could reasonably be considered inappropriate in a + professional setting + +## Our Responsibilities + +Project maintainers are responsible for clarifying the standards of acceptable +behavior and are expected to take appropriate and fair corrective action in +response to any instances of unacceptable behavior. + +Project maintainers have the right and responsibility to remove, edit, or reject +comments, commits, code, wiki edits, issues, and other contributions that are +not aligned to this Code of Conduct, or to ban temporarily or permanently any +contributor for other behaviors that they deem inappropriate, threatening, +offensive, or harmful. + +## Scope + +This Code of Conduct applies both within project spaces and in public spaces +when an individual is representing the project or its community. Examples of +representing a project or community include using an official project e-mail +address, posting via an official social media account, or acting as an appointed +representative at an online or offline event. Representation of a project may be +further defined and clarified by project maintainers. + +This Code of Conduct also applies outside the project spaces when the Project +Steward has a reasonable belief that an individual's behavior may have a +negative impact on the project or its community. + +## Conflict Resolution + +We do not believe that all conflict is bad; healthy debate and disagreement +often yield positive results. However, it is never okay to be disrespectful or +to engage in behavior that violates the project’s code of conduct. + +If you see someone violating the code of conduct, you are encouraged to address +the behavior directly with those involved. Many issues can be resolved quickly +and easily, and this gives people more control over the outcome of their +dispute. If you are unable to resolve the matter for any reason, or if the +behavior is threatening or harassing, report it. We are dedicated to providing +an environment where participants feel welcome and safe. + + +Reports should be directed to *googleapis-stewards@google.com*, the +Project Steward(s) for *Google Cloud Client Libraries*. It is the Project Steward’s duty to +receive and address reported violations of the code of conduct. They will then +work with a committee consisting of representatives from the Open Source +Programs Office and the Google Open Source Strategy team. If for any reason you +are uncomfortable reaching out to the Project Steward, please email +opensource@google.com. + +We will investigate every complaint, but you may not receive a direct response. +We will use our discretion in determining when and how to follow up on reported +incidents, which may range from not taking action to permanent expulsion from +the project and project-sponsored spaces. We will notify the accused of the +report and provide them an opportunity to discuss it before any action is taken. +The identity of the reporter will be omitted from the details of the report +supplied to the accused. In potentially harmful situations, such as ongoing +harassment or threats to anyone's safety, we may take action without notice. + +## Attribution + +This Code of Conduct is adapted from the Contributor Covenant, version 1.4, +available at +https://www.contributor-covenant.org/version/1/4/code-of-conduct.html \ No newline at end of file diff --git a/SECURITY.md b/SECURITY.md new file mode 100644 index 00000000000..8b58ae9c01a --- /dev/null +++ b/SECURITY.md @@ -0,0 +1,7 @@ +# Security Policy + +To report a security issue, please use [g.co/vulnz](https://g.co/vulnz). + +The Google Security Team will respond within 5 working days of your report on g.co/vulnz. + +We use g.co/vulnz for our intake, and do coordination and disclosure here using GitHub Security Advisory to privately discuss and fix the issue. diff --git a/bigframes/blob/_functions.py b/bigframes/blob/_functions.py new file mode 100644 index 00000000000..3dfe38811b1 --- /dev/null +++ b/bigframes/blob/_functions.py @@ -0,0 +1,728 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass +import inspect +import typing +from typing import Callable, Iterable, Union + +import google.cloud.bigquery as bigquery + +import bigframes.session +import bigframes.session._io.bigquery as bf_io_bigquery + +_PYTHON_TO_BQ_TYPES = { + int: "INT64", + float: "FLOAT64", + str: "STRING", + bytes: "BYTES", + bool: "BOOL", +} + + +@dataclass(frozen=True) +class FunctionDef: + """Definition of a Python UDF.""" + + func: Callable # function body + requirements: Iterable[str] # required packages + + +# TODO(garrettwu): migrate to bigframes UDF when it is available +class TransformFunction: + """Simple transform function class to deal with Python UDF.""" + + def __init__( + self, + func_def: FunctionDef, + session: bigframes.session.Session, + connection: str, + max_batching_rows: int, + container_cpu: Union[float, int], + container_memory: str, + ): + self._func = func_def.func + self._requirements = func_def.requirements + self._session = session + self._connection = connection + self._max_batching_rows = ( + int(max_batching_rows) if max_batching_rows > 1 else max_batching_rows + ) + self._container_cpu = container_cpu + self._container_memory = container_memory + + def _input_bq_signature(self): + sig = inspect.signature(self._func) + inputs = [] + for k, v in sig.parameters.items(): + inputs.append(f"{k} {_PYTHON_TO_BQ_TYPES[v.annotation]}") + return ", ".join(inputs) + + def _output_bq_type(self): + sig = inspect.signature(self._func) + return_annotation = sig.return_annotation + origin = typing.get_origin(return_annotation) + if origin is Union: + args = typing.get_args(return_annotation) + if len(args) == 2 and args[1] is type(None): + return _PYTHON_TO_BQ_TYPES[args[0]] + return _PYTHON_TO_BQ_TYPES[sig.return_annotation] + + def _create_udf(self): + """Create Python UDF in BQ. Return name of the UDF.""" + udf_name = str( + self._session._anon_dataset_manager.generate_unique_resource_id() + ) + + func_body = "import typing\n" + inspect.getsource(self._func) + func_name = self._func.__name__ + packages = str(list(self._requirements)) + + sql = f""" +CREATE OR REPLACE FUNCTION `{udf_name}`({self._input_bq_signature()}) +RETURNS {self._output_bq_type()} LANGUAGE python +WITH CONNECTION `{self._connection}` +OPTIONS (entry_point='{func_name}', runtime_version='python-3.11', packages={packages}, max_batching_rows={self._max_batching_rows}, container_cpu={self._container_cpu}, container_memory='{self._container_memory}') +AS r\"\"\" + + +{func_body} + + +\"\"\" + """ + + bf_io_bigquery.start_query_with_client( + self._session.bqclient, + sql, + job_config=bigquery.QueryJobConfig(), + metrics=self._session._metrics, + location=None, + project=None, + timeout=None, + query_with_job=True, + publisher=self._session._publisher, + ) + + return udf_name + + def udf(self): + """Create and return the UDF object.""" + udf_name = self._create_udf() + + # TODO(b/404605969): remove cleanups when UDF fixes dataset deletion. + self._session._function_session._update_temp_artifacts(udf_name, "") + return self._session.read_gbq_function(udf_name) + + +def exif_func(src_obj_ref_rt: str, verbose: bool) -> str: + try: + import io + import json + + from PIL import ExifTags, Image + import requests + from requests import adapters + + session = requests.Session() + session.mount("https://", adapters.HTTPAdapter(max_retries=3)) + + src_obj_ref_rt_json = json.loads(src_obj_ref_rt) + src_url = src_obj_ref_rt_json["access_urls"]["read_url"] + + response = session.get(src_url, timeout=30) + response.raise_for_status() + bts = response.content + + image = Image.open(io.BytesIO(bts)) + exif_data = image.getexif() + exif_dict = {} + + if exif_data: + for tag, value in exif_data.items(): + tag_name = ExifTags.TAGS.get(tag, tag) + # Convert non-serializable types to strings + try: + json.dumps(value) + exif_dict[tag_name] = value + except (TypeError, ValueError): + exif_dict[tag_name] = str(value) + + if verbose: + return json.dumps({"status": "", "content": json.dumps(exif_dict)}) + else: + return json.dumps(exif_dict) + + except Exception as e: + # Return error as JSON with error field + error_result = {"status": f"{type(e).__name__}: {str(e)}", "content": "{}"} + if verbose: + return json.dumps(error_result) + else: + return "{}" + + +exif_func_def = FunctionDef(exif_func, ["pillow", "requests"]) + + +# Blur images. Takes ObjectRefRuntime as JSON string. Outputs ObjectRefRuntime JSON string. +def image_blur_func( + src_obj_ref_rt: str, + dst_obj_ref_rt: str, + ksize_x: int, + ksize_y: int, + ext: str, + verbose: bool, +) -> typing.Optional[str]: + try: + import json + + import cv2 as cv # type: ignore + import numpy as np + import requests + from requests import adapters + + session = requests.Session() + session.mount("https://", adapters.HTTPAdapter(max_retries=3)) + + ext = ext or ".jpeg" + + src_obj_ref_rt_json = json.loads(src_obj_ref_rt) + dst_obj_ref_rt_json = json.loads(dst_obj_ref_rt) + + src_url = src_obj_ref_rt_json["access_urls"]["read_url"] + dst_url = dst_obj_ref_rt_json["access_urls"]["write_url"] + + response = session.get(src_url, timeout=30) + response.raise_for_status() # Raise exception for HTTP errors + bts = response.content + + nparr = np.frombuffer(bts, np.uint8) + img = cv.imdecode(nparr, cv.IMREAD_UNCHANGED) + + if img is None: + raise ValueError( + "Failed to decode image - possibly corrupted or unsupported format" + ) + + img_blurred = cv.blur(img, ksize=(ksize_x, ksize_y)) + + success, encoded = cv.imencode(ext, img_blurred) + if not success: + raise ValueError(f"Failed to encode image with extension {ext}") + + bts = encoded.tobytes() + + ext = ext.replace(".", "") + ext_mappings = {"jpg": "jpeg", "tif": "tiff"} + ext = ext_mappings.get(ext, ext) + content_type = "image/" + ext + + put_response = session.put( + url=dst_url, + data=bts, + headers={"Content-Type": content_type}, + timeout=30, + ) + put_response.raise_for_status() + + if verbose: + return json.dumps({"status": "", "content": dst_obj_ref_rt}) + else: + return dst_obj_ref_rt + + except Exception as e: + if verbose: + error_result = { + "status": f"Error: {type(e).__name__}: {str(e)}", + "content": "", + } + return json.dumps(error_result) + else: + return None + + +image_blur_def = FunctionDef(image_blur_func, ["opencv-python", "numpy", "requests"]) + + +def image_blur_to_bytes_func( + src_obj_ref_rt: str, ksize_x: int, ksize_y: int, ext: str, verbose: bool +) -> str: + import base64 + import json + + try: + import cv2 as cv # type: ignore + import numpy as np + import requests + from requests import adapters + + session = requests.Session() + session.mount("https://", adapters.HTTPAdapter(max_retries=3)) + + ext = ext or ".jpeg" + + src_obj_ref_rt_json = json.loads(src_obj_ref_rt) + src_url = src_obj_ref_rt_json["access_urls"]["read_url"] + + response = session.get(src_url, timeout=30) + response.raise_for_status() + bts = response.content + + nparr = np.frombuffer(bts, np.uint8) + img = cv.imdecode(nparr, cv.IMREAD_UNCHANGED) + if img is None: + raise ValueError( + "Failed to decode image - possibly corrupted or unsupported format" + ) + img_blurred = cv.blur(img, ksize=(ksize_x, ksize_y)) + success, encoded = cv.imencode(ext, img_blurred) + if not success: + raise ValueError(f"Failed to encode image with extension {ext}") + content = encoded.tobytes() + + encoded_content = base64.b64encode(content).decode("utf-8") + result_dict = {"status": "", "content": encoded_content} + if verbose: + return json.dumps(result_dict) + else: + return result_dict["content"] + + except Exception as e: + status = f"Error: {type(e).__name__}: {str(e)}" + encoded_content = base64.b64encode(b"").decode("utf-8") + result_dict = {"status": status, "content": encoded_content} + if verbose: + return json.dumps(result_dict) + else: + return result_dict["content"] + + +image_blur_to_bytes_def = FunctionDef( + image_blur_to_bytes_func, ["opencv-python", "numpy", "requests"] +) + + +def image_resize_func( + src_obj_ref_rt: str, + dst_obj_ref_rt: str, + dsize_x: int, + dsize_y: int, + fx: float, + fy: float, + ext: str, + verbose: bool, +) -> typing.Optional[str]: + try: + import json + + import cv2 as cv # type: ignore + import numpy as np + import requests + from requests import adapters + + session = requests.Session() + session.mount("https://", adapters.HTTPAdapter(max_retries=3)) + + ext = ext or ".jpeg" + + src_obj_ref_rt_json = json.loads(src_obj_ref_rt) + dst_obj_ref_rt_json = json.loads(dst_obj_ref_rt) + + src_url = src_obj_ref_rt_json["access_urls"]["read_url"] + dst_url = dst_obj_ref_rt_json["access_urls"]["write_url"] + + response = session.get(src_url, timeout=30) + response.raise_for_status() + bts = response.content + + nparr = np.frombuffer(bts, np.uint8) + img = cv.imdecode(nparr, cv.IMREAD_UNCHANGED) + if img is None: + raise ValueError( + "Failed to decode image - possibly corrupted or unsupported format" + ) + img_resized = cv.resize(img, dsize=(dsize_x, dsize_y), fx=fx, fy=fy) + + success, encoded = cv.imencode(ext, img_resized) + if not success: + raise ValueError(f"Failed to encode image with extension {ext}") + bts = encoded.tobytes() + + ext = ext.replace(".", "") + ext_mappings = {"jpg": "jpeg", "tif": "tiff"} + ext = ext_mappings.get(ext, ext) + content_type = "image/" + ext + + put_response = session.put( + url=dst_url, + data=bts, + headers={ + "Content-Type": content_type, + }, + timeout=30, + ) + put_response.raise_for_status() + + if verbose: + return json.dumps({"status": "", "content": dst_obj_ref_rt}) + else: + return dst_obj_ref_rt + + except Exception as e: + if verbose: + error_result = { + "status": f"Error: {type(e).__name__}: {str(e)}", + "content": "", + } + return json.dumps(error_result) + else: + return None + + +image_resize_def = FunctionDef( + image_resize_func, ["opencv-python", "numpy", "requests"] +) + + +def image_resize_to_bytes_func( + src_obj_ref_rt: str, + dsize_x: int, + dsize_y: int, + fx: float, + fy: float, + ext: str, + verbose: bool, +) -> str: + import base64 + import json + + try: + import cv2 as cv # type: ignore + import numpy as np + import requests + from requests import adapters + + session = requests.Session() + session.mount("https://", adapters.HTTPAdapter(max_retries=3)) + + ext = ext or ".jpeg" + + src_obj_ref_rt_json = json.loads(src_obj_ref_rt) + src_url = src_obj_ref_rt_json["access_urls"]["read_url"] + + response = session.get(src_url, timeout=30) + response.raise_for_status() + bts = response.content + + nparr = np.frombuffer(bts, np.uint8) + img = cv.imdecode(nparr, cv.IMREAD_UNCHANGED) + if img is None: + raise ValueError( + "Failed to decode image - possibly corrupted or unsupported format" + ) + img_resized = cv.resize(img, dsize=(dsize_x, dsize_y), fx=fx, fy=fy) + success, encoded = cv.imencode(ext, img_resized) + if not success: + raise ValueError(f"Failed to encode image with extension {ext}") + content = encoded.tobytes() + + encoded_content = base64.b64encode(content).decode("utf-8") + result_dict = {"status": "", "content": encoded_content} + if verbose: + return json.dumps(result_dict) + else: + return result_dict["content"] + + except Exception as e: + status = f"Error: {type(e).__name__}: {str(e)}" + encoded_content = base64.b64encode(b"").decode("utf-8") + result_dict = {"status": status, "content": encoded_content} + if verbose: + return json.dumps(result_dict) + else: + return result_dict["content"] + + +image_resize_to_bytes_def = FunctionDef( + image_resize_to_bytes_func, ["opencv-python", "numpy", "requests"] +) + + +def image_normalize_func( + src_obj_ref_rt: str, + dst_obj_ref_rt: str, + alpha: float, + beta: float, + norm_type: str, + ext: str, + verbose: bool, +) -> typing.Optional[str]: + try: + import json + + import cv2 as cv # type: ignore + import numpy as np + import requests + from requests import adapters + + session = requests.Session() + session.mount("https://", adapters.HTTPAdapter(max_retries=3)) + + ext = ext or ".jpeg" + + norm_type_mapping = { + "inf": cv.NORM_INF, + "l1": cv.NORM_L1, + "l2": cv.NORM_L2, + "minmax": cv.NORM_MINMAX, + } + + src_obj_ref_rt_json = json.loads(src_obj_ref_rt) + dst_obj_ref_rt_json = json.loads(dst_obj_ref_rt) + + src_url = src_obj_ref_rt_json["access_urls"]["read_url"] + dst_url = dst_obj_ref_rt_json["access_urls"]["write_url"] + + response = session.get(src_url, timeout=30) + response.raise_for_status() + bts = response.content + + nparr = np.frombuffer(bts, np.uint8) + img = cv.imdecode(nparr, cv.IMREAD_UNCHANGED) + if img is None: + raise ValueError( + "Failed to decode image - possibly corrupted or unsupported format" + ) + img_normalized = cv.normalize( + img, None, alpha=alpha, beta=beta, norm_type=norm_type_mapping[norm_type] + ) + + success, encoded = cv.imencode(ext, img_normalized) + if not success: + raise ValueError(f"Failed to encode image with extension {ext}") + bts = encoded.tobytes() + + ext = ext.replace(".", "") + ext_mappings = {"jpg": "jpeg", "tif": "tiff"} + ext = ext_mappings.get(ext, ext) + content_type = "image/" + ext + + put_response = session.put( + url=dst_url, + data=bts, + headers={ + "Content-Type": content_type, + }, + timeout=30, + ) + put_response.raise_for_status() + + if verbose: + return json.dumps({"status": "", "content": dst_obj_ref_rt}) + else: + return dst_obj_ref_rt + + except Exception as e: + if verbose: + error_result = { + "status": f"Error: {type(e).__name__}: {str(e)}", + "content": "", + } + return json.dumps(error_result) + else: + return None + + +image_normalize_def = FunctionDef( + image_normalize_func, ["opencv-python", "numpy", "requests"] +) + + +def image_normalize_to_bytes_func( + src_obj_ref_rt: str, + alpha: float, + beta: float, + norm_type: str, + ext: str, + verbose: bool, +) -> str: + import base64 + import json + + try: + import cv2 as cv # type: ignore + import numpy as np + import requests + from requests import adapters + + session = requests.Session() + session.mount("https://", adapters.HTTPAdapter(max_retries=3)) + + ext = ext or ".jpeg" + + norm_type_mapping = { + "inf": cv.NORM_INF, + "l1": cv.NORM_L1, + "l2": cv.NORM_L2, + "minmax": cv.NORM_MINMAX, + } + + src_obj_ref_rt_json = json.loads(src_obj_ref_rt) + src_url = src_obj_ref_rt_json["access_urls"]["read_url"] + + response = session.get(src_url, timeout=30) + response.raise_for_status() + bts = response.content + + nparr = np.frombuffer(bts, np.uint8) + img = cv.imdecode(nparr, cv.IMREAD_UNCHANGED) + if img is None: + raise ValueError( + "Failed to decode image - possibly corrupted or unsupported format" + ) + img_normalized = cv.normalize( + img, None, alpha=alpha, beta=beta, norm_type=norm_type_mapping[norm_type] + ) + success, encoded = cv.imencode(ext, img_normalized) + if not success: + raise ValueError(f"Failed to encode image with extension {ext}") + content = encoded.tobytes() + + encoded_content = base64.b64encode(content).decode("utf-8") + result_dict = {"status": "", "content": encoded_content} + + if verbose: + return json.dumps(result_dict) + else: + return result_dict["content"] + + except Exception as e: + status = f"Error: {type(e).__name__}: {str(e)}" + encoded_content = base64.b64encode(b"").decode("utf-8") + result_dict = {"status": status, "content": encoded_content} + if verbose: + return json.dumps(result_dict) + else: + return result_dict["content"] + + +image_normalize_to_bytes_def = FunctionDef( + image_normalize_to_bytes_func, ["opencv-python", "numpy", "requests"] +) + + +# Extracts all text from a PDF url +def pdf_extract_func(src_obj_ref_rt: str, verbose: bool) -> str: + try: + import io + import json + + from pypdf import PdfReader # type: ignore + import requests + from requests import adapters + + session = requests.Session() + session.mount("https://", adapters.HTTPAdapter(max_retries=3)) + + src_obj_ref_rt_json = json.loads(src_obj_ref_rt) + src_url = src_obj_ref_rt_json["access_urls"]["read_url"] + + response = session.get(src_url, timeout=30, stream=True) + response.raise_for_status() + pdf_bytes = response.content + + pdf_file = io.BytesIO(pdf_bytes) + reader = PdfReader(pdf_file, strict=False) + + all_text = "" + for page in reader.pages: + page_extract_text = page.extract_text() + if page_extract_text: + all_text += page_extract_text + + result_dict = {"status": "", "content": all_text} + + except Exception as e: + result_dict = {"status": str(e), "content": ""} + + if verbose: + return json.dumps(result_dict) + else: + return result_dict["content"] + + +pdf_extract_def = FunctionDef( + pdf_extract_func, ["pypdf>=5.3.1,<6.0.0", "requests", "cryptography==43.0.3"] +) + + +# Extracts text from a PDF url and chunks it simultaneously +def pdf_chunk_func( + src_obj_ref_rt: str, chunk_size: int, overlap_size: int, verbose: bool +) -> str: + try: + import io + import json + + from pypdf import PdfReader # type: ignore + import requests + from requests import adapters + + session = requests.Session() + session.mount("https://", adapters.HTTPAdapter(max_retries=3)) + + src_obj_ref_rt_json = json.loads(src_obj_ref_rt) + src_url = src_obj_ref_rt_json["access_urls"]["read_url"] + + response = session.get(src_url, timeout=30, stream=True) + response.raise_for_status() + pdf_bytes = response.content + + pdf_file = io.BytesIO(pdf_bytes) + reader = PdfReader(pdf_file, strict=False) + # extract and chunk text simultaneously + all_text_chunks = [] + curr_chunk = "" + for page in reader.pages: + page_text = page.extract_text() + if page_text: + curr_chunk += page_text + # split the accumulated text into chunks of a specific size with overlaop + # this loop implements a sliding window approach to create chunks + while len(curr_chunk) >= chunk_size: + split_idx = curr_chunk.rfind(" ", 0, chunk_size) + if split_idx == -1: + split_idx = chunk_size + actual_chunk = curr_chunk[:split_idx] + all_text_chunks.append(actual_chunk) + overlap = curr_chunk[split_idx + 1 : split_idx + 1 + overlap_size] + curr_chunk = overlap + curr_chunk[split_idx + 1 + overlap_size :] + if curr_chunk: + all_text_chunks.append(curr_chunk) + + result_dict = {"status": "", "content": all_text_chunks} + + except Exception as e: + result_dict = {"status": str(e), "content": []} + + if verbose: + return json.dumps(result_dict) + else: + return json.dumps(result_dict["content"]) + + +pdf_chunk_def = FunctionDef( + pdf_chunk_func, ["pypdf>=5.3.1,<6.0.0", "requests", "cryptography==43.0.3"] +) diff --git a/docs/README.rst b/docs/README.rst deleted file mode 100644 index a3aef5380bb..00000000000 --- a/docs/README.rst +++ /dev/null @@ -1,94 +0,0 @@ -BigQuery DataFrames (BigFrames) -=============================== - - -|GA| |pypi| |versions| - -BigQuery DataFrames (also known as BigFrames) provides a Pythonic DataFrame -and machine learning (ML) API powered by the BigQuery engine. It provides modules -for many use cases, including: - -* `bigframes.pandas `_ - is a pandas API for analytics. Many workloads can be - migrated from pandas to bigframes by just changing a few imports. -* `bigframes.ml `_ - is a scikit-learn-like API for ML. -* `bigframes.bigquery.ai `_ - are a collection of powerful AI methods, powered by Gemini. - -BigQuery DataFrames is an `open-source package `_. - -.. |GA| image:: https://img.shields.io/badge/support-GA-gold.svg - :target: https://github.com/googleapis/google-cloud-python/blob/main/README.rst#general-availability -.. |pypi| image:: https://img.shields.io/pypi/v/bigframes.svg - :target: https://pypi.org/project/bigframes/ -.. |versions| image:: https://img.shields.io/pypi/pyversions/bigframes.svg - :target: https://pypi.org/project/bigframes/ - -Getting started with BigQuery DataFrames ----------------------------------------- - -The easiest way to get started is to try the -`BigFrames quickstart `_ -in a `notebook in BigQuery Studio `_. - -To use BigFrames in your local development environment, - -1. Run ``pip install --upgrade bigframes`` to install the latest version. - -2. Setup `Application default credentials `_ - for your local development environment enviroment. - -3. Create a `GCP project with the BigQuery API enabled `_. - -4. Use the ``bigframes`` package to query data. - -.. code-block:: python - - import bigframes.pandas as bpd - - bpd.options.bigquery.project = your_gcp_project_id # Optional in BQ Studio. - bpd.options.bigquery.ordering_mode = "partial" # Recommended for performance. - df = bpd.read_gbq("bigquery-public-data.usa_names.usa_1910_2013") - print( - df.groupby("name") - .agg({"number": "sum"}) - .sort_values("number", ascending=False) - .head(10) - .to_pandas() - ) - -Documentation -------------- - -To learn more about BigQuery DataFrames, visit these pages - -* `Introduction to BigQuery DataFrames (BigFrames) `_ -* `Sample notebooks `_ -* `API reference `_ -* `Source code (GitHub) `_ - -License -------- - -BigQuery DataFrames is distributed with the `Apache-2.0 license -`_. - -It also contains code derived from the following third-party packages: - -* `Ibis `_ -* `pandas `_ -* `Python `_ -* `scikit-learn `_ -* `XGBoost `_ -* `SQLGlot `_ - -For details, see the `third_party -`_ -directory. - - -Contact Us ----------- - -For further help and provide feedback, you can email us at `bigframes-feedback@google.com `_. diff --git a/docs/README.rst b/docs/README.rst new file mode 120000 index 00000000000..89a0106941f --- /dev/null +++ b/docs/README.rst @@ -0,0 +1 @@ +../README.rst \ No newline at end of file diff --git a/docs/changelog.md b/docs/changelog.md new file mode 120000 index 00000000000..04c99a55caa --- /dev/null +++ b/docs/changelog.md @@ -0,0 +1 @@ +../CHANGELOG.md \ No newline at end of file diff --git a/renovate.json b/renovate.json new file mode 100644 index 00000000000..c7875c469bd --- /dev/null +++ b/renovate.json @@ -0,0 +1,12 @@ +{ + "extends": [ + "config:base", + "group:all", + ":preserveSemverRanges", + ":disableDependencyDashboard" + ], + "ignorePaths": [".pre-commit-config.yaml", ".kokoro/requirements.txt", "setup.py", ".github/workflows/unittest.yml"], + "pip_requirements": { + "fileMatch": ["requirements-test.txt", "samples/[\\S/]*constraints.txt", "samples/[\\S/]*constraints-test.txt"] + } +} diff --git a/samples/snippets/__init__.py b/samples/snippets/__init__.py new file mode 100644 index 00000000000..1dc90d18483 --- /dev/null +++ b/samples/snippets/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/samples/snippets/bigquery_modules_test.py b/samples/snippets/bigquery_modules_test.py new file mode 100644 index 00000000000..0cc2b1d8b5c --- /dev/null +++ b/samples/snippets/bigquery_modules_test.py @@ -0,0 +1,98 @@ +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +def test_bigquery_dataframes_examples() -> None: + # [START bigquery_dataframes_bigquery_methods_array_agg] + import bigframes.bigquery as bbq + import bigframes.pandas as bpd + + s = bpd.Series([0, 1, 2, 3, 4, 5]) + + # Group values by whether they are divisble by 2 and aggregate them into arrays + bbq.array_agg(s.groupby(s % 2 == 0)) + # False [1 3 5] + # True [0 2 4] + # dtype: list[pyarrow] + # [END bigquery_dataframes_bigquery_methods_array_agg] + + # [START bigquery_dataframes_bigquery_methods_struct] + import bigframes.bigquery as bbq + import bigframes.pandas as bpd + + # Load data from BigQuery + query_or_table = "bigquery-public-data.ml_datasets.penguins" + bq_df = bpd.read_gbq(query_or_table) + + # Create a new STRUCT Series with subfields for each column in a DataFrames. + lengths = bbq.struct( + bq_df[["culmen_length_mm", "culmen_depth_mm", "flipper_length_mm"]] + ) + + lengths.peek() + # 146 {'culmen_length_mm': 51.1, 'culmen_depth_mm': ... + # 278 {'culmen_length_mm': 48.2, 'culmen_depth_mm': ... + # 337 {'culmen_length_mm': 36.4, 'culmen_depth_mm': ... + # 154 {'culmen_length_mm': 46.5, 'culmen_depth_mm': ... + # 185 {'culmen_length_mm': 50.1, 'culmen_depth_mm': ... + # dtype: struct[pyarrow] + # [END bigquery_dataframes_bigquery_methods_struct] + + # [START bigquery_dataframes_bigquery_methods_unix_micros] + import pandas as pd + + import bigframes.bigquery as bbq + import bigframes.pandas as bpd + + # Create a series that consists of three timestamps: [1970-01-01, 1970-01-02, 1970-01-03] + s = bpd.Series(pd.date_range("1970-01-01", periods=3, freq="d", tz="UTC")) + + bbq.unix_micros(s) + # 0 0 + # 1 86400000000 + # 2 172800000000 + # dtype: Int64 + # [END bigquery_dataframes_bigquery_methods_unix_micros] + + # [START bigquery_dataframes_bigquery_methods_scalar] + import bigframes.bigquery as bbq + import bigframes.pandas as bpd + + # Load data from BigQuery + query_or_table = "bigquery-public-data.ml_datasets.penguins" + + # The sql_scalar function can be used to inject SQL syntax that is not supported + # or difficult to express with the bigframes.pandas APIs. + bq_df = bpd.read_gbq(query_or_table) + shortest = bbq.sql_scalar( + "LEAST({0}, {1}, {2})", + columns=[ + bq_df["culmen_depth_mm"], + bq_df["culmen_length_mm"], + bq_df["flipper_length_mm"], + ], + ) + + shortest.peek() + # 0 + # 149 18.9 + # 33 16.3 + # 296 17.2 + # 287 17.0 + # 307 15.0 + # dtype: Float64 + # [END bigquery_dataframes_bigquery_methods_scalar] + assert bq_df is not None + assert lengths is not None + assert shortest is not None diff --git a/samples/snippets/bqml_getting_started_test.py b/samples/snippets/bqml_getting_started_test.py new file mode 100644 index 00000000000..d249ca4ff34 --- /dev/null +++ b/samples/snippets/bqml_getting_started_test.py @@ -0,0 +1,297 @@ +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +def test_bqml_getting_started(random_model_id: str) -> None: + your_model_id = random_model_id # for example: bqml_tutorial.sample_model + + # [START bigquery_dataframes_bqml_getting_started_tutorial] + from bigframes.ml.linear_model import LogisticRegression + import bigframes.pandas as bpd + + # Start by selecting the data you'll use for training. `read_gbq` accepts + # either a SQL query or a table ID. Since this example selects from multiple + # tables via a wildcard, use SQL to define this data. Watch issue + # https://github.com/googleapis/python-bigquery-dataframes/issues/169 + # for updates to `read_gbq` to support wildcard tables. + + df = bpd.read_gbq_table( + "bigquery-public-data.google_analytics_sample.ga_sessions_*", + filters=[ + ("_table_suffix", ">=", "20160801"), + ("_table_suffix", "<=", "20170630"), + ], + ) + + # Extract the total number of transactions within + # the Google Analytics session. + # + # Because the totals column is a STRUCT data type, call + # Series.struct.field("transactions") to extract the transactions field. + # See the reference documentation below: + # https://cloud.google.com/python/docs/reference/bigframes/latest/bigframes.operations.structs.StructAccessor#bigframes_operations_structs_StructAccessor_field + transactions = df["totals"].struct.field("transactions") + + # The "label" values represent the outcome of the model's + # prediction. In this case, the model predicts if there are any + # ecommerce transactions within the Google Analytics session. + # If the number of transactions is NULL, the value in the label + # column is set to 0. Otherwise, it is set to 1. + label = transactions.notnull().map({True: 1, False: 0}).rename("label") + + # Extract the operating system of the visitor's device. + operating_system = df["device"].struct.field("operatingSystem") + operating_system = operating_system.fillna("") + + # Extract whether the visitor's device is a mobile device. + is_mobile = df["device"].struct.field("isMobile") + + # Extract the country from which the sessions originated, based on the IP address. + country = df["geoNetwork"].struct.field("country").fillna("") + + # Extract the total number of page views within the session. + pageviews = df["totals"].struct.field("pageviews").fillna(0) + + # Combine all the feature columns into a single DataFrame + # to use as training data. + features = bpd.DataFrame( + { + "os": operating_system, + "is_mobile": is_mobile, + "country": country, + "pageviews": pageviews, + } + ) + + # Logistic Regression model splits data into two classes, giving the + # a confidence score that the data is in one of the classes. + model = LogisticRegression() + model.fit(features, label) + + # The model.fit() call above created a temporary model. + # Use the to_gbq() method to write to a permanent location. + model.to_gbq( + your_model_id, # For example: "bqml_tutorial.sample_model", + replace=True, + ) + # [END bigquery_dataframes_bqml_getting_started_tutorial] + + # [START bigquery_dataframes_bqml_getting_started_tutorial_evaluate] + import bigframes.pandas as bpd + + # Select model you'll use for evaluating. `read_gbq_model` loads model data from a + # BigQuery, but you could also use the `model` object from the previous steps. + model = bpd.read_gbq_model( + your_model_id, # For example: "bqml_tutorial.sample_model", + ) + + # The filters parameter limits the number of tables scanned by the query. + # The date range scanned is July 1, 2017 to August 1, 2017. This is the + # data you're using to evaluate the predictive performance of the model. + # It was collected in the month immediately following the time period + # spanned by the training data. + df = bpd.read_gbq_table( + "bigquery-public-data.google_analytics_sample.ga_sessions_*", + filters=[ + ("_table_suffix", ">=", "20170701"), + ("_table_suffix", "<=", "20170801"), + ], + ) + + transactions = df["totals"].struct.field("transactions") + label = transactions.notnull().map({True: 1, False: 0}).rename("label") + operating_system = df["device"].struct.field("operatingSystem") + operating_system = operating_system.fillna("") + is_mobile = df["device"].struct.field("isMobile") + country = df["geoNetwork"].struct.field("country").fillna("") + pageviews = df["totals"].struct.field("pageviews").fillna(0) + features = bpd.DataFrame( + { + "os": operating_system, + "is_mobile": is_mobile, + "country": country, + "pageviews": pageviews, + } + ) + + # Some models include a convenient .score(X, y) method for evaluation with a preset accuracy metric: + + # Because you performed a logistic regression, the results include the following columns: + + # - precision — A metric for classification models. Precision identifies the frequency with + # which a model was correct when predicting the positive class. + + # - recall — A metric for classification models that answers the following question: + # Out of all the possible positive labels, how many did the model correctly identify? + + # - accuracy — Accuracy is the fraction of predictions that a classification model got right. + + # - f1_score — A measure of the accuracy of the model. The f1 score is the harmonic average of + # the precision and recall. An f1 score's best value is 1. The worst value is 0. + + # - log_loss — The loss function used in a logistic regression. This is the measure of how far the + # model's predictions are from the correct labels. + + # - roc_auc — The area under the ROC curve. This is the probability that a classifier is more confident that + # a randomly chosen positive example + # is actually positive than that a randomly chosen negative example is positive. For more information, + # see ['Classification']('https://developers.google.com/machine-learning/crash-course/classification/video-lecture') + # in the Machine Learning Crash Course. + + model.score(features, label) + # precision recall accuracy f1_score log_loss roc_auc + # 0 0.412621 0.079143 0.985074 0.132812 0.049764 0.974285 + # [1 rows x 6 columns] + # [END bigquery_dataframes_bqml_getting_started_tutorial_evaluate] + + # [START bigquery_dataframes_bqml_getting_started_tutorial_predict_by_country] + import bigframes.pandas as bpd + + # Select model you'll use for predicting. + # `read_gbq_model` loads model data from + # BigQuery, but you could also use the `model` + # object from the previous steps. + model = bpd.read_gbq_model( + your_model_id, # For example: "bqml_tutorial.sample_model", + ) + + # The filters parameter limits the number of tables scanned by the query. + # The date range scanned is July 1, 2017 to August 1, 2017. This is the + # data you're using to make the prediction. + # It was collected in the month immediately following the time period + # spanned by the training data. + df = bpd.read_gbq_table( + "bigquery-public-data.google_analytics_sample.ga_sessions_*", + filters=[ + ("_table_suffix", ">=", "20170701"), + ("_table_suffix", "<=", "20170801"), + ], + ) + + operating_system = df["device"].struct.field("operatingSystem") + operating_system = operating_system.fillna("") + is_mobile = df["device"].struct.field("isMobile") + country = df["geoNetwork"].struct.field("country").fillna("") + pageviews = df["totals"].struct.field("pageviews").fillna(0) + features = bpd.DataFrame( + { + "os": operating_system, + "is_mobile": is_mobile, + "country": country, + "pageviews": pageviews, + } + ) + # Use Logistic Regression predict method to predict results + # using your model. + # Find more information here in + # [BigFrames](https://cloud.google.com/python/docs/reference/bigframes/latest/bigframes.ml.linear_model.LogisticRegression#bigframes_ml_linear_model_LogisticRegression_predict) + + predictions = model.predict(features) + + # Call groupby method to group predicted_label by country. + # Call sum method to get the total_predicted_label by country. + total_predicted_purchases = predictions.groupby(["country"])[ + ["predicted_label"] + ].sum() + + # Call the sort_values method with the parameter + # ascending = False to get the highest values. + # Call head method to limit to the 10 highest values. + total_predicted_purchases.sort_values(ascending=False).head(10) + + # country + # United States 220 + # Taiwan 8 + # Canada 7 + # India 2 + # Japan 2 + # Turkey 2 + # Australia 1 + # Brazil 1 + # Germany 1 + # Guyana 1 + # Name: predicted_label, dtype: Int64 + + # [END bigquery_dataframes_bqml_getting_started_tutorial_predict_by_country] + + # [START bigquery_dataframes_bqml_getting_started_tutorial_predict_by_visitor] + + import bigframes.pandas as bpd + + # Select model you'll use for predicting. + # `read_gbq_model` loads model data from + # BigQuery, but you could also use the `model` + # object from the previous steps. + model = bpd.read_gbq_model( + your_model_id, # For example: "bqml_tutorial.sample_model", + ) + + # The filters parameter limits the number of tables scanned by the query. + # The date range scanned is July 1, 2017 to August 1, 2017. This is the + # data you're using to make the prediction. + # It was collected in the month immediately following the time period + # spanned by the training data. + df = bpd.read_gbq_table( + "bigquery-public-data.google_analytics_sample.ga_sessions_*", + filters=[ + ("_table_suffix", ">=", "20170701"), + ("_table_suffix", "<=", "20170801"), + ], + ) + + operating_system = df["device"].struct.field("operatingSystem") + operating_system = operating_system.fillna("") + is_mobile = df["device"].struct.field("isMobile") + country = df["geoNetwork"].struct.field("country").fillna("") + pageviews = df["totals"].struct.field("pageviews").fillna(0) + full_visitor_id = df["fullVisitorId"] + + features = bpd.DataFrame( + { + "os": operating_system, + "is_mobile": is_mobile, + "country": country, + "pageviews": pageviews, + "fullVisitorId": full_visitor_id, + } + ) + + predictions = model.predict(features) + + # Call groupby method to group predicted_label by visitor. + # Call sum method to get the total_predicted_label by visitor. + total_predicted_purchases = predictions.groupby(["fullVisitorId"])[ + ["predicted_label"] + ].sum() + + # Call the sort_values method with the parameter + # ascending = False to get the highest values. + # Call head method to limit to the 10 highest values. + total_predicted_purchases.sort_values(ascending=False).head(10) + + # fullVisitorId + # 9417857471295131045 4 + # 0376394056092189113 2 + # 0456807427403774085 2 + # 057693500927581077 2 + # 112288330928895942 2 + # 1280993661204347450 2 + # 2105122376016897629 2 + # 2158257269735455737 2 + # 2969418676126258798 2 + # 489038402765684003 2 + # Name: predicted_label, dtype: Int64 + + +# [END bigquery_dataframes_bqml_getting_started_tutorial_predict_by_visitor] diff --git a/samples/snippets/classification_boosted_tree_model_test.py b/samples/snippets/classification_boosted_tree_model_test.py new file mode 100644 index 00000000000..c79706a0200 --- /dev/null +++ b/samples/snippets/classification_boosted_tree_model_test.py @@ -0,0 +1,114 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +def test_boosted_tree_model(random_model_id: str) -> None: + your_model_id = random_model_id + # [START bigquery_dataframes_bqml_boosted_tree_prepare] + import bigframes.pandas as bpd + + input_data = bpd.read_gbq( + "bigquery-public-data.ml_datasets.census_adult_income", + columns=( + "age", + "workclass", + "marital_status", + "education_num", + "occupation", + "hours_per_week", + "income_bracket", + "functional_weight", + ), + ) + input_data["dataframe"] = bpd.Series("training", index=input_data.index,).case_when( + [ + (((input_data["functional_weight"] % 10) == 8), "evaluation"), + (((input_data["functional_weight"] % 10) == 9), "prediction"), + ] + ) + del input_data["functional_weight"] + # [END bigquery_dataframes_bqml_boosted_tree_prepare] + # [START bigquery_dataframes_bqml_boosted_tree_create] + from bigframes.ml import ensemble + + # input_data is defined in an earlier step. + training_data = input_data[input_data["dataframe"] == "training"] + X = training_data.drop(columns=["income_bracket", "dataframe"]) + y = training_data["income_bracket"] + + # create and train the model + tree_model = ensemble.XGBClassifier( + n_estimators=1, + booster="gbtree", + tree_method="hist", + max_iterations=1, # For a more accurate model, try 50 iterations. + subsample=0.85, + ) + tree_model.fit(X, y) + + tree_model.to_gbq( + your_model_id, # For example: "your-project.bqml_tutorial.tree_model" + replace=True, + ) + # [END bigquery_dataframes_bqml_boosted_tree_create] + # [START bigquery_dataframes_bqml_boosted_tree_evaluate] + # Select model you'll use for predictions. `read_gbq_model` loads model + # data from BigQuery, but you could also use the `tree_model` object + # from the previous step. + tree_model = bpd.read_gbq_model( + your_model_id, # For example: "your-project.bqml_tutorial.tree_model" + ) + + # input_data is defined in an earlier step. + evaluation_data = input_data[input_data["dataframe"] == "evaluation"] + X = evaluation_data.drop(columns=["income_bracket", "dataframe"]) + y = evaluation_data["income_bracket"] + + # The score() method evaluates how the model performs compared to the + # actual data. Output DataFrame matches that of ML.EVALUATE(). + score = tree_model.score(X, y) + score.peek() + # Output: + # precision recall accuracy f1_score log_loss roc_auc + # 0 0.671924 0.578804 0.839429 0.621897 0.344054 0.887335 + # [END bigquery_dataframes_bqml_boosted_tree_evaluate] + # [START bigquery_dataframes_bqml_boosted_tree_predict] + # Select model you'll use for predictions. `read_gbq_model` loads model + # data from BigQuery, but you could also use the `tree_model` object + # from previous steps. + tree_model = bpd.read_gbq_model( + your_model_id, # For example: "your-project.bqml_tutorial.tree_model" + ) + + # input_data is defined in an earlier step. + prediction_data = input_data[input_data["dataframe"] == "prediction"] + + predictions = tree_model.predict(prediction_data) + predictions.peek() + # Output: + # predicted_income_bracket predicted_income_bracket_probs.label predicted_income_bracket_probs.prob + # <=50K >50K 0.05183430016040802 + # <50K 0.94816571474075317 + # <=50K >50K 0.00365859130397439 + # <50K 0.99634140729904175 + # <=50K >50K 0.037775970995426178 + # <50K 0.96222406625747681 + # [END bigquery_dataframes_bqml_boosted_tree_predict] + assert input_data is not None + assert training_data is not None + assert tree_model is not None + assert evaluation_data is not None + assert score is not None + assert prediction_data is not None + assert predictions is not None diff --git a/samples/snippets/clustering_model_test.py b/samples/snippets/clustering_model_test.py new file mode 100644 index 00000000000..fec4bbcefd8 --- /dev/null +++ b/samples/snippets/clustering_model_test.py @@ -0,0 +1,35 @@ +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +def test_clustering_model() -> None: + # [START bigquery_dataframes_clustering_model] + from bigframes.ml.cluster import KMeans + import bigframes.pandas as bpd + + # Load data from BigQuery + query_or_table = "bigquery-public-data.ml_datasets.penguins" + bq_df = bpd.read_gbq(query_or_table) + + # Create the KMeans model + cluster_model = KMeans(n_clusters=10) + cluster_model.fit(bq_df["culmen_length_mm"], bq_df["sex"]) + + # Predict using the model + result = cluster_model.predict(bq_df) + # Score the model + score = cluster_model.score(bq_df) + # [END bigquery_dataframes_clustering_model] + assert result is not None + assert score is not None diff --git a/samples/snippets/conftest.py b/samples/snippets/conftest.py new file mode 100644 index 00000000000..e19cfbceb46 --- /dev/null +++ b/samples/snippets/conftest.py @@ -0,0 +1,141 @@ +# Copyright 2020 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Generator, Iterator + +from google.cloud import bigquery, storage +import pytest +import test_utils.prefixer + +import bigframes.pandas as bpd + +prefixer = test_utils.prefixer.Prefixer( + "python-bigquery-dataframes", "samples/snippets" +) + +routine_prefixer = test_utils.prefixer.Prefixer("bigframes", "") + + +@pytest.fixture(scope="session", autouse=True) +def cleanup_datasets(bigquery_client: bigquery.Client) -> None: + for dataset in bigquery_client.list_datasets(): + if prefixer.should_cleanup(dataset.dataset_id): + bigquery_client.delete_dataset( + dataset, delete_contents=True, not_found_ok=True + ) + + +@pytest.fixture(scope="session") +def bigquery_client() -> bigquery.Client: + bigquery_client = bigquery.Client() + return bigquery_client + + +@pytest.fixture(scope="session") +def storage_client(project_id: str) -> storage.Client: + return storage.Client(project=project_id) + + +@pytest.fixture(scope="session") +def project_id(bigquery_client: bigquery.Client) -> str: + return bigquery_client.project + + +@pytest.fixture(scope="session") +def gcs_bucket(storage_client: storage.Client) -> Generator[str, None, None]: + bucket_name = "bigframes_blob_test_with_data_wipeout" + + yield bucket_name + + bucket = storage_client.get_bucket(bucket_name) + for blob in bucket.list_blobs(): + blob.delete() + + +@pytest.fixture(scope="session") +def gcs_bucket_snippets(storage_client: storage.Client) -> Generator[str, None, None]: + bucket_name = "bigframes_blob_test_snippet_with_data_wipeout" + + yield bucket_name + + bucket = storage_client.get_bucket(bucket_name) + for blob in bucket.list_blobs(): + blob.delete() + + +@pytest.fixture(autouse=True) +def reset_session() -> None: + """An autouse fixture ensuring each sample runs in a fresh session. + + This allows us to have samples that query data in different locations. + """ + bpd.reset_session() + bpd.options.bigquery.location = None + + +@pytest.fixture(scope="session") +def dataset_id(bigquery_client: bigquery.Client, project_id: str) -> Iterator[str]: + dataset_id = prefixer.create_prefix() + full_dataset_id = f"{project_id}.{dataset_id}" + dataset = bigquery.Dataset(full_dataset_id) + bigquery_client.create_dataset(dataset) + yield dataset_id + bigquery_client.delete_dataset(dataset, delete_contents=True, not_found_ok=True) + + +@pytest.fixture(scope="session") +def dataset_id_eu(bigquery_client: bigquery.Client, project_id: str) -> Iterator[str]: + dataset_id = prefixer.create_prefix() + full_dataset_id = f"{project_id}.{dataset_id}" + dataset = bigquery.Dataset(full_dataset_id) + dataset.location = "EU" + bigquery_client.create_dataset(dataset) + yield dataset_id + bigquery_client.delete_dataset(dataset, delete_contents=True, not_found_ok=True) + + +@pytest.fixture +def random_model_id( + bigquery_client: bigquery.Client, project_id: str, dataset_id: str +) -> Iterator[str]: + """Create a new table ID each time, so random_model_id can be used as + target for load jobs. + """ + random_model_id = prefixer.create_prefix() + full_model_id = f"{project_id}.{dataset_id}.{random_model_id}" + yield full_model_id + bigquery_client.delete_model(full_model_id, not_found_ok=True) + + +@pytest.fixture +def random_model_id_eu( + bigquery_client: bigquery.Client, project_id: str, dataset_id_eu: str +) -> Iterator[str]: + """ + Create a new table ID each time, so random_model_id_eu can be used + as a target for load jobs. + """ + random_model_id_eu = prefixer.create_prefix() + full_model_id = f"{project_id}.{dataset_id_eu}.{random_model_id_eu}" + yield full_model_id + bigquery_client.delete_model(full_model_id, not_found_ok=True) + + +@pytest.fixture +def routine_id() -> Iterator[str]: + """Create a new BQ routine ID each time, so random_routine_id can be used as + target for udf creation. + """ + random_routine_id = routine_prefixer.create_prefix() + yield random_routine_id diff --git a/samples/snippets/create_kmeans_model_test.py b/samples/snippets/create_kmeans_model_test.py new file mode 100644 index 00000000000..7d9a43e86c2 --- /dev/null +++ b/samples/snippets/create_kmeans_model_test.py @@ -0,0 +1,162 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +def test_kmeans_sample(project_id: str, random_model_id_eu: str) -> None: + your_gcp_project_id = project_id + your_model_id = random_model_id_eu + # [START bigquery_dataframes_bqml_kmeans] + import datetime + import typing + + import pandas as pd + from shapely.geometry import Point + + import bigframes + import bigframes.bigquery as bbq + import bigframes.geopandas + import bigframes.pandas as bpd + + bigframes.options.bigquery.project = your_gcp_project_id + # Compute in the EU multi-region to query the London bicycles dataset. + bigframes.options.bigquery.location = "EU" + + # Extract the information you'll need to train the k-means model in this + # tutorial. Use the read_gbq function to represent cycle hires + # data as a DataFrame. + h = bpd.read_gbq( + "bigquery-public-data.london_bicycles.cycle_hire", + col_order=["start_station_name", "start_station_id", "start_date", "duration"], + ).rename( + columns={ + "start_station_name": "station_name", + "start_station_id": "station_id", + } + ) + + # Use GeoSeries.from_xy and BigQuery.st_distance to analyze geographical + # data. These functions determine spatial relationships between + # geographical features. + cycle_stations = bpd.read_gbq("bigquery-public-data.london_bicycles.cycle_stations") + s = bpd.DataFrame( + { + "id": cycle_stations["id"], + "xy": bigframes.geopandas.GeoSeries.from_xy( + cycle_stations["longitude"], cycle_stations["latitude"] + ), + } + ) + s_distance = bbq.st_distance(s["xy"], Point(-0.1, 51.5), use_spheroid=False) / 1000 + s = bpd.DataFrame({"id": s["id"], "distance_from_city_center": s_distance}) + + # Define Python datetime objects in the UTC timezone for range comparison, + # because BigQuery stores timestamp data in the UTC timezone. + sample_time = datetime.datetime(2015, 1, 1, 0, 0, 0, tzinfo=datetime.timezone.utc) + sample_time2 = datetime.datetime(2016, 1, 1, 0, 0, 0, tzinfo=datetime.timezone.utc) + + h = h.loc[(h["start_date"] >= sample_time) & (h["start_date"] <= sample_time2)] + + # Replace each day-of-the-week number with the corresponding "weekday" or + # "weekend" label by using the Series.map method. + h = h.assign( + isweekday=h.start_date.dt.dayofweek.map( + { + 0: "weekday", + 1: "weekday", + 2: "weekday", + 3: "weekday", + 4: "weekday", + 5: "weekend", + 6: "weekend", + } + ) + ) + + # Supplement each trip in "h" with the station distance information from + # "s" by merging the two DataFrames by station ID. + merged_df = h.merge( + right=s, + how="inner", + left_on="station_id", + right_on="id", + ) + + # Engineer features to cluster the stations. For each station, find the + # average trip duration, number of trips, and distance from city center. + stationstats = typing.cast( + bpd.DataFrame, + merged_df.groupby(["station_name", "isweekday"]).agg( + {"duration": ["mean", "count"], "distance_from_city_center": "max"} + ), + ) + stationstats.columns = pd.Index( + ["duration", "num_trips", "distance_from_city_center"] + ) + stationstats = stationstats.sort_values( + by="distance_from_city_center", ascending=True + ).reset_index() + + # Expected output results: >>> stationstats.head(3) + # station_name isweekday duration num_trips distance_from_city_center + # Borough Road... weekday 1110 5749 0.12624 + # Borough Road... weekend 2125 1774 0.12624 + # Webber Street... weekday 795 6517 0.164021 + # 3 rows × 5 columns + + # [END bigquery_dataframes_bqml_kmeans] + + # [START bigquery_dataframes_bqml_kmeans_fit] + + from bigframes.ml.cluster import KMeans + + # To determine an optimal number of clusters, construct and fit several + # K-Means objects with different values of num_clusters, find the error + # measure, and pick the point at which the error measure is at its minimum + # value. + cluster_model = KMeans(n_clusters=4) + cluster_model.fit(stationstats) + cluster_model.to_gbq( + your_model_id, # For example: "bqml_tutorial.london_station_clusters" + replace=True, + ) + # [END bigquery_dataframes_bqml_kmeans_fit] + + # [START bigquery_dataframes_bqml_kmeans_predict] + + # Select model you'll use for predictions. `read_gbq_model` loads model + # data from BigQuery, but you could also use the `cluster_model` object + # from previous steps. + cluster_model = bpd.read_gbq_model( + your_model_id, + # For example: "bqml_tutorial.london_station_clusters", + ) + + # Use 'contains' function to filter by stations containing the string + # "Kennington". + stationstats = stationstats.loc[ + stationstats["station_name"].str.contains("Kennington") + ] + + result = cluster_model.predict(stationstats) + + # Expected output results: >>>results.peek(3) + # CENTROID... NEAREST... station_name isweekday duration num_trips dist... + # 1 [{'CENTROID_ID'... Borough... weekday 1110 5749 0.13 + # 2 [{'CENTROID_ID'... Borough... weekend 2125 1774 0.13 + # 1 [{'CENTROID_ID'... Webber... weekday 795 6517 0.16 + # 3 rows × 7 columns + + # [END bigquery_dataframes_bqml_kmeans_predict] + + assert result is not None diff --git a/samples/snippets/create_multiple_timeseries_forecasting_model_test.py b/samples/snippets/create_multiple_timeseries_forecasting_model_test.py new file mode 100644 index 00000000000..0ce38e1a85d --- /dev/null +++ b/samples/snippets/create_multiple_timeseries_forecasting_model_test.py @@ -0,0 +1,175 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +def test_multiple_timeseries_forecasting_model(random_model_id: str) -> None: + your_model_id = random_model_id + + # [START bigquery_dataframes_bqml_arima_multiple_step_2_visualize] + + import bigframes.pandas as bpd + + df = bpd.read_gbq("bigquery-public-data.new_york.citibike_trips") + + features = bpd.DataFrame( + { + "num_trips": df.starttime, + "date": df["starttime"].dt.date, + } + ) + date = df["starttime"].dt.date + df.groupby([date]) + num_trips = features.groupby(["date"]).count() + + # Results from running "print(num_trips)" + + # num_trips + # date + # 2013-07-01 16650 + # 2013-07-02 22745 + # 2013-07-03 21864 + # 2013-07-04 22326 + # 2013-07-05 21842 + # 2013-07-06 20467 + # 2013-07-07 20477 + # 2013-07-08 21615 + # 2013-07-09 26641 + # 2013-07-10 25732 + # 2013-07-11 24417 + # 2013-07-12 19006 + # 2013-07-13 26119 + # 2013-07-14 29287 + # 2013-07-15 28069 + # 2013-07-16 29842 + # 2013-07-17 30550 + # 2013-07-18 28869 + # 2013-07-19 26591 + # 2013-07-20 25278 + # 2013-07-21 30297 + # 2013-07-22 25979 + # 2013-07-23 32376 + # 2013-07-24 35271 + # 2013-07-25 31084 + + num_trips.plot.line( + # Rotate the x labels so they are more visible. + rot=45, + ) + + # [END bigquery_dataframes_bqml_arima_multiple_step_2_visualize] + + # [START bigquery_dataframes_bqml_arima_multiple_step_3_fit] + from bigframes.ml import forecasting + import bigframes.pandas as bpd + + model = forecasting.ARIMAPlus( + # To reduce the query runtime with the compromise of a potential slight + # drop in model quality, you could decrease the value of the + # auto_arima_max_order. This shrinks the search space of hyperparameter + # tuning in the auto.ARIMA algorithm. + auto_arima_max_order=5, + ) + + df = bpd.read_gbq("bigquery-public-data.new_york.citibike_trips") + + # This query creates twelve time series models, one for each of the twelve + # Citi Bike start stations in the input data. If you remove this row + # filter, there would be 600+ time series to forecast. + df = df[df["start_station_name"].str.contains("Central Park")] + + features = bpd.DataFrame( + { + "start_station_name": df["start_station_name"], + "num_trips": df["starttime"], + "date": df["starttime"].dt.date, + } + ) + num_trips = features.groupby( + ["start_station_name", "date"], + as_index=False, + ).count() + + X = num_trips["date"].to_frame() + y = num_trips["num_trips"].to_frame() + + model.fit( + X, + y, + # The input data that you want to get forecasts for, + # in this case the Citi Bike station, as represented by the + # start_station_name column. + id_col=num_trips["start_station_name"].to_frame(), + ) + + # The model.fit() call above created a temporary model. + # Use the to_gbq() method to write to a permanent location. + model.to_gbq( + your_model_id, # For example: "bqml_tutorial.nyc_citibike_arima_model", + replace=True, + ) + # [END bigquery_dataframes_bqml_arima_multiple_step_3_fit] + + # [START bigquery_dataframes_bqml_arima_multiple_step_4_evaluate] + # Evaluate the time series models by using the summary() function. The summary() + # function shows you the evaluation metrics of all the candidate models evaluated + # during the process of automatic hyperparameter tuning. + summary = model.summary() + print(summary.peek()) + + # Expected output: + # start_station_name non_seasonal_p non_seasonal_d non_seasonal_q has_drift log_likelihood AIC variance ... + # 1 Central Park West & W 72 St 0 1 5 False -1966.449243 3944.898487 1215.689281 ... + # 8 Central Park W & W 96 St 0 0 5 False -274.459923 562.919847 655.776577 ... + # 9 Central Park West & W 102 St 0 0 0 False -226.639918 457.279835 258.83582 ... + # 11 Central Park West & W 76 St 1 1 2 False -1700.456924 3408.913848 383.254161 ... + # 4 Grand Army Plaza & Central Park S 0 1 5 False -5507.553498 11027.106996 624.138741 ... + # [END bigquery_dataframes_bqml_arima_multiple_step_4_evaluate] + + # [START bigquery_dataframes_bqml_arima_multiple_step_5_coefficients] + coef = model.coef_ + print(coef.peek()) + + # Expected output: + # start_station_name ar_coefficients ma_coefficients intercept_or_drift + # 5 Central Park West & W 68 St [] [-0.41014089 0.21979212 -0.59854213 -0.251438... 0.0 + # 6 Central Park S & 6 Ave [] [-0.71488957 -0.36835772 0.61008532 0.183290... 0.0 + # 0 Central Park West & W 85 St [] [-0.39270166 -0.74494638 0.76432596 0.489146... 0.0 + # 3 W 82 St & Central Park West [-0.50219511 -0.64820817] [-0.20665325 0.67683137 -0.68108631] 0.0 + # 11 W 106 St & Central Park West [-0.70442887 -0.66885553 -0.25030325 -0.34160669] [] 0.0 + # [END bigquery_dataframes_bqml_arima_multiple_step_5_coefficients] + + # [START bigquery_dataframes_bqml_arima_multiple_step_6_forecast] + prediction = model.predict(horizon=3, confidence_level=0.9) + + print(prediction.peek()) + # Expected output: + # forecast_timestamp start_station_name forecast_value standard_error confidence_level ... + # 4 2016-10-01 00:00:00+00:00 Central Park S & 6 Ave 302.377201 32.572948 0.9 ... + # 14 2016-10-02 00:00:00+00:00 Central Park North & Adam Clayton Powell Blvd 263.917567 45.284082 0.9 ... + # 1 2016-09-25 00:00:00+00:00 Central Park West & W 85 St 189.574706 39.874856 0.9 ... + # 20 2016-10-02 00:00:00+00:00 Central Park West & W 72 St 175.474862 40.940794 0.9 ... + # 12 2016-10-01 00:00:00+00:00 W 106 St & Central Park West 63.88163 18.088868 0.9 ... + # [END bigquery_dataframes_bqml_arima_multiple_step_6_forecast] + # [START bigquery_dataframes_bqml_arima_multiple_step_7_explain] + explain = model.predict_explain(horizon=3, confidence_level=0.9) + + print(explain.peek(5)) + # Expected output: + # time_series_timestamp start_station_name time_series_type time_series_data time_series_adjusted_data standard_error confidence_level prediction_interval_lower_bound prediction_interval_upper_bound trend seasonal_period_yearly seasonal_period_quarterly seasonal_period_monthly seasonal_period_weekly seasonal_period_daily holiday_effect spikes_and_dips step_changes residual + # 0 2013-07-01 00:00:00+00:00 Central Park S & 6 Ave history 69.0 154.168527 32.572948 0.0 35.477484 -28.402102 0.0 -85.168527 147.093145 + # 1 2013-07-01 00:00:00+00:00 Grand Army Plaza & Central Park S history 79.0 79.0 24.982769 0.0 43.46428 -30.01599 0.0 0.0 65.55171 + # 2 2013-07-02 00:00:00+00:00 Central Park S & 6 Ave history 180.0 204.045651 32.572948 147.093045 72.498327 -15.545721 0.0 -85.168527 61.122876 + # 3 2013-07-02 00:00:00+00:00 Grand Army Plaza & Central Park S history 129.0 99.556269 24.982769 65.551665 45.836432 -11.831828 0.0 0.0 29.443731 + # 4 2013-07-03 00:00:00+00:00 Central Park S & 6 Ave history 115.0 205.968236 32.572948 191.32754 59.220766 -44.580071 0.0 -85.168527 -5.799709 + # [END bigquery_dataframes_bqml_arima_multiple_step_7_explain] diff --git a/samples/snippets/create_single_timeseries_forecasting_model_test.py b/samples/snippets/create_single_timeseries_forecasting_model_test.py new file mode 100644 index 00000000000..9965da28176 --- /dev/null +++ b/samples/snippets/create_single_timeseries_forecasting_model_test.py @@ -0,0 +1,127 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (t +# you may not use this file except in compliance wi +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in +# distributed under the License is distributed on a +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, eit +# See the License for the specific language governi +# limitations under the License. + + +def test_create_single_timeseries() -> None: + # [START bigquery_dataframes_single_timeseries_forecasting_model_tutorial] + import bigframes.pandas as bpd + + # Start by loading the historical data from BigQuerythat you want to analyze and forecast. + # This clause indicates that you are querying the ga_sessions_* tables in the google_analytics_sample dataset. + # Read and visualize the time series you want to forecast. + df = bpd.read_gbq("bigquery-public-data.google_analytics_sample.ga_sessions_*") + parsed_date = bpd.to_datetime(df.date, format="%Y%m%d", utc=True) + parsed_date.name = "parsed_date" + visits = df["totals"].struct.field("visits") + visits.name = "total_visits" + total_visits = visits.groupby(parsed_date).sum() + + # Expected output: total_visits.head() + # parsed_date + # 2016-08-01 00:00:00+00:00 1711 + # 2016-08-02 00:00:00+00:00 2140 + # 2016-08-03 00:00:00+00:00 2890 + # 2016-08-04 00:00:00+00:00 3161 + # 2016-08-05 00:00:00+00:00 2702 + # Name: total_visits, dtype: Int64 + + total_visits.plot.line() + + # [END bigquery_dataframes_single_timeseries_forecasting_model_tutorial] + + # [START bigquery_dataframes_single_timeseries_forecasting_model_tutorial_create] + from bigframes.ml import forecasting + import bigframes.pandas as bpd + + # Create a time series model to forecast total site visits: + # The auto_arima option defaults to True, so the auto.ARIMA algorithm automatically + # tunes the hyperparameters in the model. + # The data_frequency option defaults to 'auto_frequency so the training + # process automatically infers the data frequency of the input time series. + # The decompose_time_series option defaults to True, so that information about + # the time series data is returned when you evaluate the model in the next step. + model = forecasting.ARIMAPlus() + model.auto_arima = True + model.data_frequency = "auto_frequency" + model.decompose_time_series = True + + # Use the data loaded in the previous step to fit the model + training_data = total_visits.to_frame().reset_index(drop=False) + + X = training_data[["parsed_date"]] + y = training_data[["total_visits"]] + + model.fit(X, y) + # [END bigquery_dataframes_single_timeseries_forecasting_model_tutorial_create] + + # [START bigquery_dataframes_single_timeseries_forecasting_model_tutorial_coef] + coef = model.coef_ + print(coef.peek()) + + # Expected output: + # ar_coefficients ma_coefficients intercept_or_drift + # 0 [0.40944762] [-0.81168198] 0.0 + # [END bigquery_dataframes_single_timeseries_forecasting_model_tutorial_coef] + + # [START bigquery_dataframes_single_timeseries_forecasting_model_tutorial_evaluate] + # Evaluate the time series models by using the summary() function. The summary() + # function shows you the evaluation metrics of all the candidate models evaluated + # during the process of automatic hyperparameter tuning. + summary = model.summary( + show_all_candidate_models=True, + ) + print(summary.peek()) + + # Expected output: + # row non_seasonal_p non_seasonal_d non_seasonal_q has_drift log_likelihood AIC variance seasonal_periods has_holiday_effect has_spikes_and_dips has_step_changes error_message + # 0 0 1 3 True -2464.255656 4938.511313 42772.506055 ['WEEKLY'] False False True + # 1 2 1 0 False -2473.141651 4952.283303 44942.416463 ['WEEKLY'] False False True + # 2 1 1 0 False -2479.880885 4963.76177 46642.953433 ['WEEKLY'] False False True + # 3 0 1 1 False -2470.632377 4945.264753 44319.379307 ['WEEKLY'] False False True + # 4 2 1 1 True -2463.671247 4937.342493 42633.299513 ['WEEKLY'] False False True + # [END bigquery_dataframes_single_timeseries_forecasting_model_tutorial_evaluate] + + # [START bigquery_dataframes_single_timeseries_forecasting_model_tutorial_forecast] + prediction = model.predict(horizon=30, confidence_level=0.8) + + print(prediction.peek()) + # Expected output: + # forecast_timestamp forecast_value standard_error confidence_level prediction_interval_lower_bound prediction_interval_upper_bound confidence_interval_lower_bound confidence_interval_upper_bound + # 11 2017-08-13 00:00:00+00:00 1845.439732 328.060405 0.8 1424.772257 2266.107208 1424.772257 2266.107208 + # 29 2017-08-31 00:00:00+00:00 2615.993932 431.286628 0.8 2062.960849 3169.027015 2062.960849 3169.027015 + # 7 2017-08-09 00:00:00+00:00 2639.285993 300.301186 0.8 2254.213792 3024.358193 2254.213792 3024.358193 + # 25 2017-08-27 00:00:00+00:00 1853.735689 410.596551 0.8 1327.233216 2380.238162 1327.233216 2380.238162 + # 1 2017-08-03 00:00:00+00:00 2621.33159 241.093355 0.8 2312.180802 2930.482379 2312.180802 2930.482379 + # [END bigquery_dataframes_single_timeseries_forecasting_model_tutorial_forecast] + + # [START bigquery_dataframes_single_timeseries_forecasting_model_tutorial_explain_forecast] + ex_pred = model.predict_explain(horizon=30, confidence_level=0.8) + + print(ex_pred.head(4)) + # Expected output: + # time_series_timestamp time_series_type time_series_data time_series_adjusted_data standard_error confidence_level prediction_interval_lower_bound prediction_interval_upper_bound trend seasonal_period_yearly seasonal_period_quarterly seasonal_period_monthly seasonal_period_weekly seasonal_period_daily holiday_effect spikes_and_dips step_changes residual + # 0 2016-08-01 00:00:00+00:00 history 1711.0 505.716474 206.939556 0.0 169.611938 1205.283526 336.104536 + # 1 2016-08-02 00:00:00+00:00 history 2140.0 623.137701 206.939556 336.104428 287.033273 1205.283526 311.578773 + # 2 2016-08-03 00:00:00+00:00 history 2890.0 1008.655091 206.939556 563.514213 445.140878 1205.283526 676.061383 + # 3 2016-08-04 00:00:00+00:00 history 3161.0 1389.40959 206.939556 986.317236 403.092354 1205.283526 566.306884 + # 4 2016-08-05 00:00:00+00:00 history 2702.0 1394.395741 206.939556 1248.707386 145.688355 1205.283526 102.320733 + # 5 2016-08-06 00:00:00+00:00 history 1663.0 437.09243 206.939556 1188.59004 -751.49761 1205.283526 20.624044 + # [END bigquery_dataframes_single_timeseries_forecasting_model_tutorial_explain_forecast] + assert coef is not None + assert ex_pred is not None + assert summary is not None + assert model is not None + assert parsed_date is not None + assert prediction is not None + assert total_visits is not None diff --git a/samples/snippets/data_visualization_test.py b/samples/snippets/data_visualization_test.py new file mode 100644 index 00000000000..64cbbe05112 --- /dev/null +++ b/samples/snippets/data_visualization_test.py @@ -0,0 +1,149 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (t +# you may not use this file except in compliance wi +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in +# distributed under the License is distributed on a +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, eit +# See the License for the specific language governi +# limitations under the License. + + +def test_data_visualization() -> None: + # [START bigquery_dataframes_data_visualization_penguin_histogram] + import bigframes.pandas as bpd + + penguins = bpd.read_gbq("bigquery-public-data.ml_datasets.penguins") + penguins["culmen_depth_mm"].plot.hist(bins=40) + # [END bigquery_dataframes_data_visualization_penguin_histogram] + + # [START bigquery_dataframes_data_visualization_noaa_line_chart] + import bigframes.pandas as bpd + + noaa_surface = bpd.read_gbq("bigquery-public-data.noaa_gsod.gsod2021") + + # Calculate median temperature for each day + noaa_surface_median_temps = noaa_surface[["date", "temp"]].groupby("date").median() + + noaa_surface_median_temps.plot.line() + # [END bigquery_dataframes_data_visualization_noaa_line_chart] + + # [START bigquery_dataframes_data_visualization_usa_names_area_chart] + import bigframes.pandas as bpd + + usa_names = bpd.read_gbq("bigquery-public-data.usa_names.usa_1910_2013") + + # Count the occurences of the target names each year. The result is a dataframe with a multi-index. + name_counts = ( + usa_names[usa_names["name"].isin(("Mary", "Emily", "Lisa"))] + .groupby(("year", "name"))["number"] + .sum() + ) + + # Flatten the index of the dataframe so that the counts for each name has their own columns. + name_counts = name_counts.unstack(level=1).fillna(0) + + name_counts.plot.area(stacked=False, alpha=0.5) + # [END bigquery_dataframes_data_visualization_usa_names_area_chart] + + # [START bigquery_dataframes_data_visualization_penguin_bar_chart] + import bigframes.pandas as bpd + + penguins = bpd.read_gbq("bigquery-public-data.ml_datasets.penguins") + + penguin_count_by_sex = ( + penguins[penguins["sex"].isin(("MALE", "FEMALE"))] + .groupby("sex")["species"] + .count() + ) + penguin_count_by_sex.plot.bar() + # [END bigquery_dataframes_data_visualization_penguin_bar_chart] + + # [START bigquery_dataframes_data_visualization_taxi_scatter_plot] + import bigframes.pandas as bpd + + taxi_trips = bpd.read_gbq( + "bigquery-public-data.new_york_taxi_trips.tlc_yellow_trips_2021" + ).dropna() + + # Data Cleaning + taxi_trips = taxi_trips[ + taxi_trips["trip_distance"].between(0, 10, inclusive="right") + ] + taxi_trips = taxi_trips[taxi_trips["fare_amount"].between(0, 50, inclusive="right")] + + # If you are using partial ordering mode, you will also need to assign an order to your dataset. + # Otherwise, the next line can be skipped. + taxi_trips = taxi_trips.sort_values("pickup_datetime") + + taxi_trips.plot.scatter(x="trip_distance", y="fare_amount", alpha=0.5) + # [END bigquery_dataframes_data_visualization_taxi_scatter_plot] + + # [START bigquery_dataframes_data_visualization_noaa_sampling_n] + import bigframes.pandas as bpd + + noaa_surface = bpd.read_gbq("bigquery-public-data.noaa_gsod.gsod2021") + + # Calculate median temperature for each day + noaa_surface_median_temps = noaa_surface[["date", "temp"]].groupby("date").median() + + noaa_surface_median_temps.plot.line(sampling_n=40) + # [END bigquery_dataframes_data_visualization_noaa_sampling_n] + + # [START bigquery_dataframes_data_visualization_usa_names_subplots] + import bigframes.pandas as bpd + + usa_names = bpd.read_gbq("bigquery-public-data.usa_names.usa_1910_2013") + + # Count the occurences of the target names each year. The result is a dataframe with a multi-index. + name_counts = ( + usa_names[usa_names["name"].isin(("Mary", "Emily", "Lisa"))] + .groupby(("year", "name"))["number"] + .sum() + ) + + # Flatten the index of the dataframe so that the counts for each name has their own columns. + name_counts = name_counts.unstack(level=1).fillna(0) + + name_counts.plot.area(subplots=True, alpha=0.5) + # [END bigquery_dataframes_data_visualization_usa_names_subplots] + + # [START bigquery_dataframes_data_visualization_taxi_scatter_multidimension] + import bigframes.pandas as bpd + + taxi_trips = bpd.read_gbq( + "bigquery-public-data.new_york_taxi_trips.tlc_yellow_trips_2021" + ).dropna() + + # Data Cleaning + taxi_trips = taxi_trips[ + taxi_trips["trip_distance"].between(0, 10, inclusive="right") + ] + taxi_trips = taxi_trips[taxi_trips["fare_amount"].between(0, 50, inclusive="right")] + + # If you are using partial ordering mode, you also need to assign an order to your dataset. + # Otherwise, the next line can be skipped. + taxi_trips = taxi_trips.sort_values("pickup_datetime") + + taxi_trips["passenger_count_scaled"] = taxi_trips["passenger_count"] * 30 + + taxi_trips.plot.scatter( + x="trip_distance", + xlabel="trip distance (miles)", + y="fare_amount", + ylabel="fare amount (usd)", + alpha=0.5, + s="passenger_count_scaled", + label="passenger_count", + c="tip_amount", + cmap="jet", + colorbar=True, + legend=True, + figsize=(15, 7), + sampling_n=1000, + ) + # [END bigquery_dataframes_data_visualization_taxi_scatter_multidimension] diff --git a/samples/snippets/explore_query_result_test.py b/samples/snippets/explore_query_result_test.py new file mode 100644 index 00000000000..7d4b241e4cc --- /dev/null +++ b/samples/snippets/explore_query_result_test.py @@ -0,0 +1,70 @@ +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +def test_bigquery_dataframes_explore_query_result() -> None: + # [START bigquery_dataframes_explore_query_result] + import bigframes.pandas as bpd + + # Load data from BigQuery + query_or_table = "bigquery-public-data.ml_datasets.penguins" + bq_df = bpd.read_gbq(query_or_table) + + # Inspect one of the columns (or series) of the DataFrame: + bq_df["body_mass_g"] + + # Compute the mean of this series: + average_body_mass = bq_df["body_mass_g"].mean() + print(f"average_body_mass: {average_body_mass}") + + # Find the heaviest species using the groupby operation to calculate the + # mean body_mass_g: + ( + bq_df["body_mass_g"] + .groupby(by=bq_df["species"]) + .mean() + .sort_values(ascending=False) + .head(10) + ) + + # Create the Linear Regression model + from bigframes.ml.linear_model import LinearRegression + + # Filter down to the data we want to analyze + adelie_data = bq_df[bq_df.species == "Adelie Penguin (Pygoscelis adeliae)"] + + # Drop the columns we don't care about + adelie_data = adelie_data.drop(columns=["species"]) + + # Drop rows with nulls to get our training data + training_data = adelie_data.dropna() + + # Pick feature columns and label column + X = training_data[ + [ + "island", + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + "sex", + ] + ] + y = training_data[["body_mass_g"]] + + model = LinearRegression(fit_intercept=False) + model.fit(X, y) + model.score(X, y) + # [END bigquery_dataframes_explore_query_result] + assert average_body_mass is not None + assert model is not None diff --git a/samples/snippets/gemini_model_test.py b/samples/snippets/gemini_model_test.py new file mode 100644 index 00000000000..fe5d7d5b1e7 --- /dev/null +++ b/samples/snippets/gemini_model_test.py @@ -0,0 +1,46 @@ +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +def test_gemini_text_generator_model() -> None: + # Determine project id, in this case prefer the one set in the environment + # variable GOOGLE_CLOUD_PROJECT (if any) + import os + + PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT", "bigframes-dev") + REGION = "us" + CONN_NAME = "bigframes-default-connection" + + # [START bigquery_dataframes_gemini_model] + from bigframes.ml.llm import GeminiTextGenerator + import bigframes.pandas as bpd + + # Create the Gemini LLM model + session = bpd.get_global_session() + connection = f"{PROJECT_ID}.{REGION}.{CONN_NAME}" + model = GeminiTextGenerator( + session=session, connection_name=connection, model_name="gemini-2.0-flash-001" + ) + + df_api = bpd.read_csv("gs://cloud-samples-data/vertex-ai/bigframe/df.csv") + + # Prepare the prompts and send them to the LLM model for prediction + df_prompt_prefix = "Generate Pandas sample code for DataFrame." + df_prompt = df_prompt_prefix + df_api["API"] + + # Predict using the model + df_pred = model.predict(df_prompt.to_frame(), max_output_tokens=1024) + # [END bigquery_dataframes_gemini_model] + assert df_pred["ml_generate_text_llm_result"] is not None + assert df_pred["ml_generate_text_llm_result"].iloc[0] is not None diff --git a/samples/snippets/imported_onnx_model_test.py b/samples/snippets/imported_onnx_model_test.py new file mode 100644 index 00000000000..87157ee60d3 --- /dev/null +++ b/samples/snippets/imported_onnx_model_test.py @@ -0,0 +1,43 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (t +# you may not use this file except in compliance wi +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in +# distributed under the License is distributed on a +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, eit +# See the License for the specific language governi +# limitations under the License. + + +def test_imported_sklearn_onnx_model() -> None: + # Determine project id, in this case prefer the one set in the environment + # variable GOOGLE_CLOUD_PROJECT (if any) + import os + + PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT", "bigframes-dev") + + # [START bigquery_dataframes_imported_sklearn_onnx_tutorial_import_onnx_models] + import bigframes + from bigframes.ml.imported import ONNXModel + + bigframes.options.bigquery.project = PROJECT_ID + # You can change the location to one of the valid locations: https://cloud.google.com/bigquery/docs/locations#supported_locations + bigframes.options.bigquery.location = "US" + + imported_onnx_model = ONNXModel( + model_path="gs://cloud-samples-data/bigquery/ml/onnx/pipeline_rf.onnx" + ) + # [END bigquery_dataframes_imported_sklearn_onnx_tutorial_import_onnx_models] + assert imported_onnx_model is not None + + # [START bigquery_dataframes_imported_sklearn_onnx_tutorial_make_predictions] + import bigframes.pandas as bpd + + df = bpd.read_gbq("bigquery-public-data.ml_datasets.iris") + predictions = imported_onnx_model.predict(df) + predictions.peek(5) + # [END bigquery_dataframes_imported_sklearn_onnx_tutorial_make_predictions] diff --git a/samples/snippets/imported_tensorflow_model_test.py b/samples/snippets/imported_tensorflow_model_test.py new file mode 100644 index 00000000000..7dcf0ffe6d2 --- /dev/null +++ b/samples/snippets/imported_tensorflow_model_test.py @@ -0,0 +1,44 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (t +# you may not use this file except in compliance wi +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in +# distributed under the License is distributed on a +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, eit +# See the License for the specific language governi +# limitations under the License. + + +def test_imported_tensorflow_model() -> None: + # Determine project id, in this case prefer the one set in the environment + # variable GOOGLE_CLOUD_PROJECT (if any) + import os + + PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT", "bigframes-dev") + + # [START bigquery_dataframes_imported_tensorflow_tutorial_import_tensorflow_models] + import bigframes + from bigframes.ml.imported import TensorFlowModel + + bigframes.options.bigquery.project = PROJECT_ID + # You can change the location to one of the valid locations: https://cloud.google.com/bigquery/docs/locations#supported_locations + bigframes.options.bigquery.location = "US" + + imported_tensorflow_model = TensorFlowModel( + model_path="gs://cloud-training-demos/txtclass/export/exporter/1549825580/*" + ) + # [END bigquery_dataframes_imported_tensorflow_tutorial_import_tensorflow_models] + assert imported_tensorflow_model is not None + + # [START bigquery_dataframes_imported_tensorflow_tutorial_make_predictions] + import bigframes.pandas as bpd + + df = bpd.read_gbq("bigquery-public-data.hacker_news.full") + df_pred = df.rename(columns={"title": "input"}) + predictions = imported_tensorflow_model.predict(df_pred) + predictions.head(5) + # [END bigquery_dataframes_imported_tensorflow_tutorial_make_predictions] diff --git a/samples/snippets/limit_single_timeseries_forecasting_model_test.py b/samples/snippets/limit_single_timeseries_forecasting_model_test.py new file mode 100644 index 00000000000..6a9f14e3835 --- /dev/null +++ b/samples/snippets/limit_single_timeseries_forecasting_model_test.py @@ -0,0 +1,64 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (t +# you may not use this file except in compliance wi +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in +# distributed under the License is distributed on a +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, eit +# See the License for the specific language governi +# limitations under the License. + + +def test_limit_single_timeseries(random_model_id: str) -> None: + your_model_id = random_model_id + + # [START bigquery_dataframes_bqml_limit_forecast_visualize] + import bigframes.pandas as bpd + + df = bpd.read_gbq("bigquery-public-data.new_york.citibike_trips") + + features = bpd.DataFrame( + { + "num_trips": df.starttime, + "date": df["starttime"].dt.date, + } + ) + num_trips = features.groupby(["date"]).count() + + num_trips.plot.line() + # [END bigquery_dataframes_bqml_limit_forecast_visualize] + + # [START bigquery_dataframes_bqml_limit_forecast_create] + from bigframes.ml import forecasting + import bigframes.pandas as bpd + + df = bpd.read_gbq("bigquery-public-data.new_york.citibike_trips") + + features = bpd.DataFrame( + { + "start_station_id": df["start_station_id"], + "num_trips": df.starttime, + "date": df["starttime"].dt.date, + } + ) + num_trips = features.groupby(["date", "start_station_id"], as_index=False).count() + model = forecasting.ARIMAPlus() + + X = num_trips[["date"]] + y = num_trips[["num_trips"]] + id_col = num_trips[["start_station_id"]] + + model.fit(X, y, id_col=id_col) + + model.to_gbq( + your_model_id, # For example: "bqml_tutorial.nyc_citibike_arima_model", + replace=True, + ) + # [END bigquery_dataframes_bqml_limit_forecast_create] + assert df is not None + assert features is not None + assert num_trips is not None diff --git a/samples/snippets/linear_regression_tutorial_test.py b/samples/snippets/linear_regression_tutorial_test.py new file mode 100644 index 00000000000..8fc1c5ad61b --- /dev/null +++ b/samples/snippets/linear_regression_tutorial_test.py @@ -0,0 +1,125 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +def test_linear_regression(random_model_id: str) -> None: + your_model_id = random_model_id + # [START bigquery_dataframes_bqml_linear_regression] + from bigframes.ml.linear_model import LinearRegression + import bigframes.pandas as bpd + + # Load data from BigQuery + bq_df = bpd.read_gbq("bigquery-public-data.ml_datasets.penguins") + + # Drop rows with nulls to get training data + training_data = bq_df.dropna(subset=["body_mass_g"]) + + # Specify your feature (or input) columns and the label (or output) column: + feature_columns = training_data.drop(columns=["body_mass_g"]) + label_columns = training_data[["body_mass_g"]] + + # Create the linear model + model = LinearRegression() + model.fit(feature_columns, label_columns) + model.to_gbq( + your_model_id, # For example: "bqml_tutorial.penguins_model" + replace=True, + ) + # [END bigquery_dataframes_bqml_linear_regression] + # [START bigquery_dataframes_bqml_linear_evaluate] + import bigframes.pandas as bpd + + # Select the model you will be evaluating. `read_gbq_model` loads model data from + # BigQuery, but you could also use the `model` object from the previous steps. + model = bpd.read_gbq_model( + your_model_id, # For example: "bqml_tutorial.penguins_model" + ) + + # Score the model with input data defined in an earlier step to compare + # model predictions on feature_columns to true labels in label_columns. + score = model.score(feature_columns, label_columns) + # Expected output results: + # index mean_absolute_error mean_squared_error mean_squared_log_error median_absolute_error r2_score explained_variance + # 0 227.012237 81838.159892 0.00507 173.080816 0.872377 0.872377 + # 1 rows x 6 columns + # [END bigquery_dataframes_bqml_linear_evaluate] + # [START bigquery_dataframes_bqml_linear_predict] + # Select the model you'll use for predictions. `read_gbq_model` loads + # model data from BigQuery, but you could also use the `model` object + # object from previous steps. + model = bpd.read_gbq_model( + your_model_id, + # For example: "bqml_tutorial.penguins_model", + ) + + # Load data from BigQuery + bq_df = bpd.read_gbq("bigquery-public-data.ml_datasets.penguins") + + # Use 'contains' function to filter by island containing the string + # "Biscoe". + biscoe_data = bq_df.loc[bq_df["island"].str.contains("Biscoe")] + + result = model.predict(biscoe_data) + + # Expected output results: + # predicted_body_mass_g species island culmen_length_mm culmen_depth_mm body_mass_g flipper_length_mm sex + # 23 4681.782896 Gentoo penguin (Pygoscelis papua) Biscoe + # 332 4740.7907 Gentoo penguin (Pygoscelis papua) Biscoe 46.2 14.4 214.0 4650.0 + # 160 4731.310452 Gentoo penguin (Pygoscelis papua) Biscoe 44.5 14.3 216.0 4100.0 + # [END bigquery_dataframes_bqml_linear_predict] + # [START bigquery_dataframes_bqml_linear_predict_explain] + # Use 'predict_explain' function to understand why the model is generating these prediction results. + # 'predict_explain'is an extended version of the 'predict' function that not only outputs prediction results, but also outputs additional columns to explain the prediction results. + # Using the trained model and utilizing data specific to Biscoe Island, explain the predictions of the top 3 features + explained = model.predict_explain(biscoe_data, top_k_features=3) + + # Expected results: + # predicted_body_mass_g top_feature_attributions baseline_prediction_value prediction_value approximation_error species island culmen_length_mm culmen_depth_mm flipper_length_mm body_mass_g sex + # 0 5413.510134 [{'feature': 'island', 'attribution': 7348.877... -5320.222128 5413.510134 0.0 Gentoo penguin (Pygoscelis papua) Biscoe 45.2 16.4 223.0 5950.0 MALE + # 1 4768.351092 [{'feature': 'island', 'attribution': 7348.877... -5320.222128 4768.351092 0.0 Gentoo penguin (Pygoscelis papua) Biscoe 46.5 14.5 213.0 4400.0 FEMALE + # 2 3235.896372 [{'feature': 'island', 'attribution': 7348.877... -5320.222128 3235.896372 0.0 Adelie Penguin (Pygoscelis adeliae) Biscoe 37.7 16.0 183.0 3075.0 FEMALE + # 3 5349.603734 [{'feature': 'island', 'attribution': 7348.877... -5320.222128 5349.603734 0.0 Gentoo penguin (Pygoscelis papua) Biscoe 46.4 15.6 221.0 5000.0 MALE + # 4 4637.165037 [{'feature': 'island', 'attribution': 7348.877... -5320.222128 4637.165037 0.0 Gentoo penguin (Pygoscelis papua) Biscoe 46.1 13.2 211.0 4500.0 FEMALE + # [END bigquery_dataframes_bqml_linear_predict_explain] + # [START bigquery_dataframes_bqml_linear_global_explain] + # To use the `global_explain()` function, the model must be recreated with `enable_global_explain` set to `True`. + model = LinearRegression(enable_global_explain=True) + + # The model must the be fitted before it can be saved to BigQuery and then explained. + training_data = bq_df.dropna(subset=["body_mass_g"]) + X = training_data.drop(columns=["body_mass_g"]) + y = training_data[["body_mass_g"]] + model.fit(X, y) + model.to_gbq("bqml_tutorial.penguins_model", replace=True) + + # Explain the model + explain_model = model.global_explain() + + # Expected results: + # attribution + # feature + # island 5737.315921 + # species 4073.280549 + # sex 622.070896 + # flipper_length_mm 193.612051 + # culmen_depth_mm 117.084944 + # culmen_length_mm 94.366793 + # [END bigquery_dataframes_bqml_linear_global_explain] + assert explain_model is not None + assert feature_columns is not None + assert label_columns is not None + assert model is not None + assert score is not None + assert result is not None + assert explained is not None diff --git a/samples/snippets/load_data_from_bigquery_test.py b/samples/snippets/load_data_from_bigquery_test.py new file mode 100644 index 00000000000..4523eece97f --- /dev/null +++ b/samples/snippets/load_data_from_bigquery_test.py @@ -0,0 +1,24 @@ +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +def test_bigquery_dataframes_load_data_from_bigquery() -> None: + # [START bigquery_dataframes_load_data_from_bigquery] + # Create a DataFrame from a BigQuery table: + import bigframes.pandas as bpd + + query_or_table = "bigquery-public-data.ml_datasets.penguins" + bq_df = bpd.read_gbq(query_or_table) + # [END bigquery_dataframes_load_data_from_bigquery] + assert bq_df is not None diff --git a/samples/snippets/load_data_from_biquery_job_test.py b/samples/snippets/load_data_from_biquery_job_test.py new file mode 100644 index 00000000000..4f1ddc062fd --- /dev/null +++ b/samples/snippets/load_data_from_biquery_job_test.py @@ -0,0 +1,56 @@ +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +def test_bigquery_dataframes_load_data_from_bigquery_job() -> None: + # Determine project id, in this case prefer the one set in the environment + # variable GOOGLE_CLOUD_PROJECT (if any) + import os + + your_project_id = os.getenv("GOOGLE_CLOUD_PROJECT", "bigframes-dev") + + # Construct a BigQuery client object. + from google.cloud import bigquery + + client = bigquery.Client(project=your_project_id, location="us") + + query = """ + SELECT * + FROM `bigquery-public-data.ml_datasets.penguins` + LIMIT 20 + """ + query_job = client.query(query) + JOB_ID = query_job.job_id + + # [START bigquery_dataframes_load_data_from_bigquery_job] + from google.cloud import bigquery + + import bigframes.pandas as bpd + + # Project ID inserted based on the query results selected to explore + project = your_project_id + # Location inserted based on the query results selected to explore + location = "us" + client = bigquery.Client(project=project, location=location) + + # Job ID inserted based on the query results selcted to explore + job_id = JOB_ID + job = client.get_job(job_id) + destination = str(job.destination) + + # Load data from a BigQuery table using BigFrames DataFrames: + bq_df = bpd.read_gbq_table(destination) + + # [END bigquery_dataframes_load_data_from_bigquery_job] + assert bq_df is not None diff --git a/samples/snippets/load_data_from_csv_test.py b/samples/snippets/load_data_from_csv_test.py new file mode 100644 index 00000000000..cc96b92fb8c --- /dev/null +++ b/samples/snippets/load_data_from_csv_test.py @@ -0,0 +1,25 @@ +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +def test_bigquery_dataframes_load_data_from_csv() -> None: + # [START bigquery_dataframes_load_data_from_csv] + import bigframes.pandas as bpd + + filepath_or_buffer = "gs://cloud-samples-data/bigquery/us-states/us-states.csv" + df_from_gcs = bpd.read_csv(filepath_or_buffer) + # Display the first few rows of the DataFrame: + df_from_gcs.head() + # [END bigquery_dataframes_load_data_from_csv] + assert df_from_gcs is not None diff --git a/samples/snippets/logistic_regression_prediction_test.py b/samples/snippets/logistic_regression_prediction_test.py new file mode 100644 index 00000000000..dd92f8f3e3e --- /dev/null +++ b/samples/snippets/logistic_regression_prediction_test.py @@ -0,0 +1,151 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""BigQuery DataFrames code samples for +https://cloud.google.com/bigquery/docs/logistic-regression-prediction. +""" + + +def test_logistic_regression_prediction(random_model_id: str) -> None: + your_model_id = random_model_id + + # [START bigquery_dataframes_logistic_regression_prediction_examine] + import bigframes.pandas as bpd + + df = bpd.read_gbq( + "bigquery-public-data.ml_datasets.census_adult_income", + columns=( + "age", + "workclass", + "marital_status", + "education_num", + "occupation", + "hours_per_week", + "income_bracket", + "functional_weight", + ), + max_results=100, + ) + df.peek() + # Output: + # age workclass marital_status education_num occupation hours_per_week income_bracket functional_weight + # 47 Local-gov Married-civ-spouse 13 Prof-specialty 40 >50K 198660 + # 56 Private Never-married 9 Adm-clerical 40 <=50K 85018 + # 40 Private Married-civ-spouse 12 Tech-support 40 >50K 285787 + # 34 Self-emp-inc Married-civ-spouse 9 Craft-repair 54 >50K 207668 + # 23 Private Married-civ-spouse 10 Handlers-cleaners 40 <=50K 40060 + # [END bigquery_dataframes_logistic_regression_prediction_examine] + + # [START bigquery_dataframes_logistic_regression_prediction_prepare] + import bigframes.pandas as bpd + + input_data = bpd.read_gbq( + "bigquery-public-data.ml_datasets.census_adult_income", + columns=( + "age", + "workclass", + "marital_status", + "education_num", + "occupation", + "hours_per_week", + "income_bracket", + "functional_weight", + ), + ) + input_data["dataframe"] = bpd.Series("training", index=input_data.index,).case_when( + [ + (((input_data["functional_weight"] % 10) == 8), "evaluation"), + (((input_data["functional_weight"] % 10) == 9), "prediction"), + ] + ) + del input_data["functional_weight"] + # [END bigquery_dataframes_logistic_regression_prediction_prepare] + + # [START bigquery_dataframes_logistic_regression_prediction_create_model] + import bigframes.ml.linear_model + + # input_data is defined in an earlier step. + training_data = input_data[input_data["dataframe"] == "training"] + X = training_data.drop(columns=["income_bracket", "dataframe"]) + y = training_data["income_bracket"] + + census_model = bigframes.ml.linear_model.LogisticRegression( + # Balance the class labels in the training data by setting + # class_weight="balanced". + # + # By default, the training data is unweighted. If the labels + # in the training data are imbalanced, the model may learn to + # predict the most popular class of labels more heavily. In + # this case, most of the respondents in the dataset are in the + # lower income bracket. This may lead to a model that predicts + # the lower income bracket too heavily. Class weights balance + # the class labels by calculating the weights for each class in + # inverse proportion to the frequency of that class. + class_weight="balanced", + max_iterations=15, + ) + census_model.fit(X, y) + + census_model.to_gbq( + your_model_id, # For example: "your-project.census.census_model" + replace=True, + ) + # [END bigquery_dataframes_logistic_regression_prediction_create_model] + + # [START bigquery_dataframes_logistic_regression_prediction_evaluate_model] + # Select model you'll use for predictions. `read_gbq_model` loads model + # data from BigQuery, but you could also use the `census_model` object + # from previous steps. + census_model = bpd.read_gbq_model( + your_model_id, # For example: "your-project.census.census_model" + ) + + # input_data is defined in an earlier step. + evaluation_data = input_data[input_data["dataframe"] == "evaluation"] + X = evaluation_data.drop(columns=["income_bracket", "dataframe"]) + y = evaluation_data["income_bracket"] + + # The score() method evaluates how the model performs compared to the + # actual data. Output DataFrame matches that of ML.EVALUATE(). + score = census_model.score(X, y) + score.peek() + # Output: + # precision recall accuracy f1_score log_loss roc_auc + # 0 0.685764 0.536685 0.83819 0.602134 0.350417 0.882953 + # [END bigquery_dataframes_logistic_regression_prediction_evaluate_model] + + # [START bigquery_dataframes_logistic_regression_prediction_predict_income_bracket] + # Select model you'll use for predictions. `read_gbq_model` loads model + # data from BigQuery, but you could also use the `census_model` object + # from previous steps. + census_model = bpd.read_gbq_model( + your_model_id, # For example: "your-project.census.census_model" + ) + + # input_data is defined in an earlier step. + prediction_data = input_data[input_data["dataframe"] == "prediction"] + + predictions = census_model.predict(prediction_data) + predictions.peek() + # Output: + # predicted_income_bracket predicted_income_bracket_probs age workclass ... occupation hours_per_week income_bracket dataframe + # 18004 <=50K [{'label': ' >50K', 'prob': 0.0763305999358786... 75 ? ... ? 6 <=50K prediction + # 18886 <=50K [{'label': ' >50K', 'prob': 0.0448866871906495... 73 ? ... ? 22 >50K prediction + # 31024 <=50K [{'label': ' >50K', 'prob': 0.0362982319421936... 69 ? ... ? 1 <=50K prediction + # 31022 <=50K [{'label': ' >50K', 'prob': 0.0787836112058324... 75 ? ... ? 5 <=50K prediction + # 23295 <=50K [{'label': ' >50K', 'prob': 0.3385373037905673... 78 ? ... ? 32 <=50K prediction + # [END bigquery_dataframes_logistic_regression_prediction_predict_income_bracket] + + # TODO(tswast): Implement ML.EXPLAIN_PREDICT() and corresponding sample. + # TODO(tswast): Implement ML.GLOBAL_EXPLAIN() and corresponding sample. diff --git a/samples/snippets/mf_explicit_model_test.py b/samples/snippets/mf_explicit_model_test.py new file mode 100644 index 00000000000..fb54b7271c3 --- /dev/null +++ b/samples/snippets/mf_explicit_model_test.py @@ -0,0 +1,162 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (t +# you may not use this file except in compliance wi +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in +# distributed under the License is distributed on a +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, eit +# See the License for the specific language governi +# limitations under the License. + + +def test_explicit_matrix_factorization(random_model_id: str) -> None: + your_model_id = random_model_id + + # [START bigquery_dataframes_bqml_mf_explicit_create_dataset] + import google.cloud.bigquery + + bqclient = google.cloud.bigquery.Client() + bqclient.create_dataset("bqml_tutorial", exists_ok=True) + # [END bigquery_dataframes_bqml_mf_explicit_create_dataset] + + # [START bigquery_dataframes_bqml_mf_explicit_upload_movielens] + import io + import zipfile + + import google.api_core.exceptions + import requests + + try: + # Check if you've already created the Movielens tables to avoid downloading + # and uploading the dataset unnecessarily. + bqclient.get_table("bqml_tutorial.ratings") + bqclient.get_table("bqml_tutorial.movies") + except google.api_core.exceptions.NotFound: + # Download the https://grouplens.org/datasets/movielens/1m/ dataset. + ml1m = requests.get("http://files.grouplens.org/datasets/movielens/ml-1m.zip") + ml1m_file = io.BytesIO(ml1m.content) + ml1m_zip = zipfile.ZipFile(ml1m_file) + + # Upload the ratings data into the ratings table. + with ml1m_zip.open("ml-1m/ratings.dat") as ratings_file: + ratings_content = ratings_file.read() + + ratings_csv = io.BytesIO(ratings_content.replace(b"::", b",")) + ratings_config = google.cloud.bigquery.LoadJobConfig() + ratings_config.source_format = "CSV" + ratings_config.write_disposition = "WRITE_TRUNCATE" + ratings_config.schema = [ + google.cloud.bigquery.SchemaField("user_id", "INT64"), + google.cloud.bigquery.SchemaField("item_id", "INT64"), + google.cloud.bigquery.SchemaField("rating", "FLOAT64"), + google.cloud.bigquery.SchemaField("timestamp", "TIMESTAMP"), + ] + bqclient.load_table_from_file( + ratings_csv, "bqml_tutorial.ratings", job_config=ratings_config + ).result() + + # Upload the movie data into the movies table. + with ml1m_zip.open("ml-1m/movies.dat") as movies_file: + movies_content = movies_file.read() + + movies_csv = io.BytesIO(movies_content.replace(b"::", b"@")) + movies_config = google.cloud.bigquery.LoadJobConfig() + movies_config.source_format = "CSV" + movies_config.field_delimiter = "@" + movies_config.write_disposition = "WRITE_TRUNCATE" + movies_config.schema = [ + google.cloud.bigquery.SchemaField("movie_id", "INT64"), + google.cloud.bigquery.SchemaField("movie_title", "STRING"), + google.cloud.bigquery.SchemaField("genre", "STRING"), + ] + bqclient.load_table_from_file( + movies_csv, "bqml_tutorial.movies", job_config=movies_config + ).result() + # [END bigquery_dataframes_bqml_mf_explicit_upload_movielens] + + # [START bigquery_dataframes_bqml_mf_explicit_create] + from bigframes.ml import decomposition + import bigframes.pandas as bpd + + # Load data from BigQuery + bq_df = bpd.read_gbq( + "bqml_tutorial.ratings", columns=("user_id", "item_id", "rating") + ) + + # Create the Matrix Factorization model + model = decomposition.MatrixFactorization( + num_factors=34, + feedback_type="explicit", + user_col="user_id", + item_col="item_id", + rating_col="rating", + l2_reg=9.83, + ) + model.fit(bq_df) + model.to_gbq( + your_model_id, replace=True # For example: "bqml_tutorial.mf_explicit" + ) + # [END bigquery_dataframes_bqml_mf_explicit_create] + # [START bigquery_dataframes_bqml_mf_explicit_evaluate] + # Evaluate the model using the score() function + model.score(bq_df) + # Output: + # mean_absolute_error mean_squared_error mean_squared_log_error median_absolute_error r2_score explained_variance + # 0.485403 0.395052 0.025515 0.390573 0.68343 0.68343 + # [END bigquery_dataframes_bqml_mf_explicit_evaluate] + # [START bigquery_dataframes_bqml_mf_explicit_recommend_df] + # Use predict() to get the predicted rating for each movie for 5 users + subset = bq_df[["user_id"]].head(5) + predicted = model.predict(subset) + print(predicted) + # Output: + # predicted_rating user_id item_id rating + # 0 4.206146 4354 968 4.0 + # 1 4.853099 3622 3521 5.0 + # 2 2.679067 5543 920 2.0 + # 3 4.323458 445 3175 5.0 + # 4 3.476911 5535 235 4.0 + # [END bigquery_dataframes_bqml_mf_explicit_recommend_df] + # [START bigquery_dataframes_bqml_mf_explicit_recommend_model] + # import bigframes.bigquery as bbq + + # Load movies + movies = bpd.read_gbq("bqml_tutorial.movies") + + # Merge the movies df with the previously created predicted df + merged_df = bpd.merge(predicted, movies, left_on="item_id", right_on="movie_id") + + # Separate users and predicted data, setting the index to 'movie_id' + users = merged_df[["user_id", "movie_id"]].set_index("movie_id") + + # Take the predicted data and sort it in descending order by 'predicted_rating', setting the index to 'movie_id' + sort_data = ( + merged_df[["movie_title", "genre", "predicted_rating", "movie_id"]] + .sort_values(by="predicted_rating", ascending=False) + .set_index("movie_id") + ) + + # re-merge the separated dfs by index + merged_user = sort_data.join(users, how="outer") + + # group the users and set the user_id as the index + merged_user.groupby("user_id").head(5).set_index("user_id").sort_index() + print(merged_user) + # Output: + # movie_title genre predicted_rating + # user_id + # 1 Saving Private Ryan (1998) Action|Drama|War 5.19326 + # 1 Fargo (1996) Crime|Drama|Thriller 4.996954 + # 1 Driving Miss Daisy (1989) Drama 4.983671 + # 1 Ben-Hur (1959) Action|Adventure|Drama 4.877622 + # 1 Schindler's List (1993) Drama|War 4.802336 + # 2 Saving Private Ryan (1998) Action|Drama|War 5.19326 + # 2 Braveheart (1995) Action|Drama|War 5.174145 + # 2 Gladiator (2000) Action|Drama 5.066372 + # 2 On Golden Pond (1981) Drama 5.01198 + # 2 Driving Miss Daisy (1989) Drama 4.983671 + # [END bigquery_dataframes_bqml_mf_explicit_recommend_model] diff --git a/samples/snippets/multimodal_test.py b/samples/snippets/multimodal_test.py new file mode 100644 index 00000000000..ce04d511346 --- /dev/null +++ b/samples/snippets/multimodal_test.py @@ -0,0 +1,526 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +def test_multimodal_dataframe(gcs_bucket_snippets: str) -> None: + # destination folder must be in a GCS bucket that the BQ connection service account (default or user provided) has write access to. + dst_bucket = f"gs://{gcs_bucket_snippets}" + # [START bigquery_dataframes_multimodal_dataframe_create] + import bigframes + + # Flags to control preview image/video preview size + bigframes.options.display.blob_display_width = 300 + + import bigframes.pandas as bpd + + # Create blob columns from wildcard path. + df_image = bpd.from_glob_path( + "gs://cloud-samples-data/bigquery/tutorials/cymbal-pets/images/*", name="image" + ) + # Other ways are: from string uri column + # df = bpd.DataFrame({"uri": ["gs:///", "gs:///"]}) + # df["blob_col"] = df["uri"].str.to_blob() + + # From an existing object table + # df = bpd.read_gbq_object_table("", name="blob_col") + + # Take only the 5 images to deal with. Preview the content of the Mutimodal DataFrame + df_image = df_image.head(5) + df_image + # [END bigquery_dataframes_multimodal_dataframe_create] + + # [START bigquery_dataframes_multimodal_dataframe_merge] + # Combine unstructured data with structured data + df_image["author"] = ["alice", "bob", "bob", "alice", "bob"] # type: ignore + df_image["content_type"] = df_image["image"].blob.content_type() + df_image["size"] = df_image["image"].blob.size() + df_image["updated"] = df_image["image"].blob.updated() + df_image + # [END bigquery_dataframes_multimodal_dataframe_merge] + + # [START bigquery_dataframes_multimodal_dataframe_filter] + # Filter images and display, you can also display audio and video types. Use width/height parameters to constrain window sizes. + df_image[df_image["author"] == "alice"]["image"].blob.display() + # [END bigquery_dataframes_multimodal_dataframe_filter] + + # [START bigquery_dataframes_multimodal_dataframe_image_transform] + df_image["blurred"] = df_image["image"].blob.image_blur( + (20, 20), dst=f"{dst_bucket}/image_blur_transformed/", engine="opencv" + ) + df_image["resized"] = df_image["image"].blob.image_resize( + (300, 200), dst=f"{dst_bucket}/image_resize_transformed/", engine="opencv" + ) + df_image["normalized"] = df_image["image"].blob.image_normalize( + alpha=50.0, + beta=150.0, + norm_type="minmax", + dst=f"{dst_bucket}/image_normalize_transformed/", + engine="opencv", + ) + + # You can also chain functions together + df_image["blur_resized"] = df_image["blurred"].blob.image_resize( + (300, 200), dst=f"{dst_bucket}/image_blur_resize_transformed/", engine="opencv" + ) + df_image + # [END bigquery_dataframes_multimodal_dataframe_image_transform] + + # [START bigquery_dataframes_multimodal_dataframe_ml_text] + from bigframes.ml import llm + + gemini = llm.GeminiTextGenerator(model_name="gemini-2.0-flash-001") + + # Deal with first 2 images as example + df_image = df_image.head(2) + + # Ask the same question on the images + df_image = df_image.head(2) + answer = gemini.predict(df_image, prompt=["what item is it?", df_image["image"]]) + answer[["ml_generate_text_llm_result", "image"]] + # [END bigquery_dataframes_multimodal_dataframe_ml_text] + + # [START bigquery_dataframes_multimodal_dataframe_ml_text_alt] + # Ask different questions + df_image["question"] = [ # type: ignore + "what item is it?", + "what color is the picture?", + ] + answer_alt = gemini.predict( + df_image, prompt=[df_image["question"], df_image["image"]] + ) + answer_alt[["ml_generate_text_llm_result", "image"]] + # [END bigquery_dataframes_multimodal_dataframe_ml_text_alt] + + # [START bigquery_dataframes_multimodal_dataframe_ml_embed] + # Generate embeddings on images + embed_model = llm.MultimodalEmbeddingGenerator() + embeddings = embed_model.predict(df_image["image"]) + embeddings + # [END bigquery_dataframes_multimodal_dataframe_ml_embed] + + # [START bigquery_dataframes_multimodal_dataframe_pdf_chunk] + # PDF chunking + df_pdf = bpd.from_glob_path( + "gs://cloud-samples-data/bigquery/tutorials/cymbal-pets/documents/*", name="pdf" + ) + df_pdf["chunked"] = df_pdf["pdf"].blob.pdf_chunk(engine="pypdf") + chunked = df_pdf["chunked"].explode() + chunked + # [END bigquery_dataframes_multimodal_dataframe_pdf_chunk] + assert df_image is not None + assert answer is not None + assert answer_alt is not None + assert embeddings is not None + assert chunked is not None + + +def test_multimodal_example(gcs_bucket_snippets: str) -> None: + BUCKET = gcs_bucket_snippets + # [START bigquery_dataframes_multimodal_load] + import bigframes.bigquery as bbq + import bigframes.pandas as bpd + + bbq.load_data( + "cymbal_pets.products", + write_disposition="OVERWRITE", + from_files_options={ + "format": "avro", + "uris": [ + "gs://cloud-samples-data/bigquery/tutorials/cymbal-pets/tables/products/products_*.avro" + ], + }, + ) + # [END bigquery_dataframes_multimodal_load] + + # [START bigquery_dataframes_multimodal_create_images] + bbq.create_external_table( + "cymbal_pets.product_images", + replace=True, + connection_name="us.cymbal_conn", + options={ + "object_metadata": "SIMPLE", + "uris": [ + "gs://cloud-samples-data/bigquery/tutorials/cymbal-pets/images/*.png" + ], + }, + ) + # [END bigquery_dataframes_multimodal_create_images] + + # [START bigquery_dataframes_multimodal_create_manuals] + bbq.create_external_table( + "cymbal_pets.product_manuals", + replace=True, + connection_name="us.cymbal_conn", + options={ + "object_metadata": "SIMPLE", + "uris": [ + "gs://cloud-samples-data/bigquery/tutorials/cymbal-pets/documents/*.pdf" + ], + }, + ) + # [END bigquery_dataframes_multimodal_create_manuals] + + # [START bigquery_dataframes_multimodal_create_gemini] + gemini_model = bbq.ml.create_model( + "cymbal_pets.gemini", + replace=True, + connection_name="us.cymbal_conn", + options={"endpoint": "gemini-2.5-flash"}, + ) + # [END bigquery_dataframes_multimodal_create_gemini] + + # [START bigquery_dataframes_multimodal_create_embedding] + embedding_model = bbq.ml.create_model( + "cymbal_pets.embedding_model", + replace=True, + connection_name="us.cymbal_conn", + options={"endpoint": "multimodalembedding@001"}, + ) + # [END bigquery_dataframes_multimodal_create_embedding] + + # [START bigquery_dataframes_multimodal_create_df_products_mm] + df_images = bpd.read_gbq("SELECT * FROM cymbal_pets.product_images") + df_products = bpd.read_gbq("cymbal_pets.products") + + df_products_mm = df_images.merge(df_products, on="uri").drop(columns="uri") + df_products_mm = df_products_mm.rename(columns={"ref": "image"}) + # [END bigquery_dataframes_multimodal_create_df_products_mm] + + # [START bigquery_dataframes_multimodal_show_df_products_mm] + df_products_mm[["product_name", "image"]] + # [END bigquery_dataframes_multimodal_show_df_products_mm] + + # [START bigquery_dataframes_multimodal_image_description] + df_products_mm["url"] = bbq.obj.get_access_url( + df_products_mm["image"], "R" + ).to_frame() + df_products_mm["prompt0"] = "Can you describe the following image?" + + df_products_mm["prompt"] = bbq.struct(df_products_mm[["prompt0", "url"]]) + df_products_mm = bbq.ai.generate_table( + gemini_model, df_products_mm, output_schema={"image_description": "STRING"} + ) + + df_products_mm = df_products_mm[ + [ + "product_id", + "product_name", + "brand", + "category", + "subcategory", + "animal_type", + "search_keywords", + "price", + "description", + "inventory_level", + "supplier_id", + "average_rating", + "image", + "image_description", + ] + ] + # [END bigquery_dataframes_multimodal_image_description] + + # [START bigquery_dataframes_multimodal_generate_animal_type] + df_prompt = bbq.obj.get_access_url(df_products_mm["image"], "R").to_frame() + df_prompt[ + "prompt0" + ] = "For the image of a pet product, concisely generate the following metadata: 1) animal_type and 2) 5 SEO search keywords, and 3) product subcategory." + + df_products_mm["prompt"] = bbq.struct(df_prompt[["prompt0", "image"]]) + + df_products_mm = df_products_mm.drop( + columns=["animal_type", "search_keywords", "subcategory"] + ) + df_products_mm = bbq.ai.generate_table( + gemini_model, + df_products_mm, + output_schema="animal_type STRING, search_keywords ARRAY, subcategory STRING", + ) + # [END bigquery_dataframes_multimodal_generate_animal_type] + + # [START bigquery_dataframes_multimodal_show_animal_type] + df_products_mm[ + [ + "product_name", + "image_description", + "animal_type", + "search_keywords", + "subcategory", + ] + ] + # [END bigquery_dataframes_multimodal_show_animal_type] + + # [START bigquery_dataframes_multimodal_brand_description] + df_agg = df_products_mm[ + ["image", "description", "category", "subcategory", "brand"] + ] + df_agg["image"] = bbq.obj.get_access_url(df_products_mm["image"], "R") + df_agg = bbq.array_agg(df_agg.groupby(by=["brand"])) + + df_agg["cnt"] = bbq.array_length(df_agg["image"]) + + df_prompt = df_agg[["image", "description", "category", "subcategory"]] + df_prompt[ + "prompt0" + ] = "Use the images and text to give one concise brand description for a website brand page. Return the description only. " + + df_agg["prompt"] = bbq.struct( + df_prompt[["prompt0", "image", "description", "category", "subcategory"]] + ) + + df_agg = df_agg.reset_index() + + df_agg = bbq.ai.generate_table( + gemini_model, df_agg, output_schema={"brand_description": "STRING"} + ) + df_agg[["brand", "brand_description", "cnt"]] + # [END bigquery_dataframes_multimodal_brand_description] + + # [START bigquery_dataframes_multimodal_define_to_grayscale] + @bpd.udf( + dataset="cymbal_pets", + name="to_grayscale", + packages=["numpy", "opencv-python"], + bigquery_connection="us.cymbal_conn", + max_batching_rows=1, + ) + def to_grayscale(src_ref: str, dst_ref: str) -> str: + import json + from urllib.request import Request, urlopen + + import cv2 as cv + import numpy as np + + src_json = json.loads(src_ref) + srcUrl = src_json["access_urls"]["read_url"] + + dst_json = json.loads(dst_ref) + dstUrl = dst_json["access_urls"]["write_url"] + + req = urlopen(srcUrl) + arr = np.asarray(bytearray(req.read()), dtype=np.uint8) + img = cv.imdecode(arr, -1) # 'Load it as it is' + + # Convert the image to grayscale + gray_image = cv.cvtColor(img, cv.COLOR_BGR2GRAY) + + # Send POST request to the URL + _, img_encoded = cv.imencode(".png", gray_image) + + req = Request( + url=dstUrl, + data=img_encoded.tobytes(), + method="PUT", + headers={ + "Content-Type": "image/png", + }, + ) + with urlopen(req): + pass + return dst_ref + + # [END bigquery_dataframes_multimodal_define_to_grayscale] + + # [START bigquery_dataframes_multimodal_apply_to_grayscale] + df_grayscale = df_products_mm[["product_id", "product_name", "image"]] + df_grayscale[ + "gray_image_uri" + ] = f"gs://{BUCKET}/cymbal-pets-images/grayscale/" + df_grayscale[ + "image" + ].struct.field( + "uri" + ).str.extract( + r"([^/]+)$" + ) + + df_grayscale["gray_image"] = bbq.obj.make_ref( + df_grayscale["gray_image_uri"], "us.cymbal_conn" + ) + + df_grayscale["image_url"] = bbq.to_json_string( + bbq.obj.get_access_url(df_grayscale["image"], "r") + ) + df_grayscale["gray_image_url"] = bbq.to_json_string( + bbq.obj.get_access_url(df_grayscale["gray_image"], "rw") + ) + + df_grayscale[["image_url", "gray_image_url"]].apply(to_grayscale, axis=1) + # [END bigquery_dataframes_multimodal_apply_to_grayscale] + + # [START bigquery_dataframes_multimodal_define_chunk_pdf] + @bpd.udf( + dataset="cymbal_pets", + name="chunk_pdf", + packages=["pypdf"], + bigquery_connection="us.cymbal_conn", + max_batching_rows=1, + ) + def chunk_pdf(src_ref: str, chunk_size: int, overlap_size: int) -> list[str]: + import io + import json + from urllib.request import urlopen + + from pypdf import PdfReader # type: ignore + + src_json = json.loads(src_ref) + srcUrl = src_json["access_urls"]["read_url"] + + req = urlopen(srcUrl) + pdf_file = io.BytesIO(bytearray(req.read())) + reader = PdfReader(pdf_file, strict=False) + + # extract and chunk text simultaneously + all_text_chunks = [] + curr_chunk = "" + for page in reader.pages: + page_text = page.extract_text() + if page_text: + curr_chunk += page_text + # split the accumulated text into chunks of a specific size with overlaop + # this loop implements a sliding window approach to create chunks + while len(curr_chunk) >= chunk_size: + split_idx = curr_chunk.rfind(" ", 0, chunk_size) + if split_idx == -1: + split_idx = chunk_size + actual_chunk = curr_chunk[:split_idx] + all_text_chunks.append(actual_chunk) + overlap = curr_chunk[split_idx + 1 : split_idx + 1 + overlap_size] + curr_chunk = overlap + curr_chunk[split_idx + 1 + overlap_size :] + if curr_chunk: + all_text_chunks.append(curr_chunk) + + return all_text_chunks + + # [END bigquery_dataframes_multimodal_define_chunk_pdf] + + # [START bigquery_dataframes_multimodal_apply_chunk_pdf] + df_manuals = bpd.read_gbq("SELECT * FROM cymbal_pets.product_manuals") + df_manuals["url"] = bbq.to_json_string( + bbq.obj.get_access_url(df_manuals["ref"], "R") + ) + + df_manuals["chunk_size"] = 1000 + df_manuals["overlap_size"] = 100 + + df_manuals["chunked"] = df_manuals[["url", "chunk_size", "overlap_size"]].apply( + chunk_pdf, axis=1 + ) + # [END bigquery_dataframes_multimodal_apply_chunk_pdf] + + # [START bigquery_dataframes_multimodal_analyze_pdf] + df_chunked = df_manuals["chunked"].explode().to_frame() + df_chunked[ + "prompt0" + ] = "Can you summarize the product manual as bullet points? Highlight the legal clauses" + + df_chunked["prompt"] = bbq.struct(df_chunked[["prompt0", "chunked"]]) + + result = bbq.ai.generate_text(gemini_model, df_chunked["prompt"]) + result + # [END bigquery_dataframes_multimodal_analyze_pdf] + + # [START bigquery_dataframes_multimodal_create_embed_table] + df_products_mm["content"] = bbq.obj.get_access_url(df_products_mm["image"], "R") + df_embed = bbq.ai.generate_embedding( + embedding_model, df_products_mm[["content", "product_id"]] + ) + + df_embed.to_gbq("cymbal_pets.products_embedding", if_exists="replace") + # [END bigquery_dataframes_multimodal_create_embed_table] + + # [START bigquery_dataframes_multimodal_vector_search] + df_image = bpd.DataFrame( + { + "uri": [ + "gs://cloud-samples-data/bigquery/tutorials/cymbal-pets/images/cozy-naps-cat-scratching-post-with-condo.png" + ] + } + ).cache() + df_image["image"] = bbq.obj.make_ref(df_image["uri"], "us.cymbal_conn") + df_search = bbq.ai.generate_embedding( + embedding_model, + bbq.obj.get_access_url(bbq.obj.fetch_metadata(df_image["image"]), "R"), + ) + + search_result = bbq.vector_search( + "cymbal_pets.products_embedding", "embedding", df_search["embedding"] + ) + search_result + # [END bigquery_dataframes_multimodal_vector_search] + + # [START bigquery_dataframes_create_external_table_all] + bbq.create_external_table( + "cymbal_pets.product_manuals_all", + replace=True, + connection_name="us.cymbal_conn", + options={ + "object_metadata": "SIMPLE", + "uris": [ + "gs://cloud-samples-data/bigquery/tutorials/cymbal-pets/documents/*.pdf", + "gs://cloud-samples-data/bigquery/tutorials/cymbal-pets/document_chunks/*.pdf", + ], + }, + ) + # [END bigquery_dataframes_create_external_table_all] + + # [START bigquery_dataframes_create_manual_to_chunks] + df1 = bpd.read_gbq("SELECT * FROM cymbal_pets.product_manuals_all").sort_values( + "uri" + ) + df2 = df1.copy() + df1["name"] = df1["uri"].str.extract(r".*/([^.]*).[^/]+") + df2["name"] = df2["uri"].str.extract(r".*/([^.]*)_page[0-9]+.[^/]+") + df_manuals_all = df1.merge(df2, on="name") + df_manuals_agg = ( + bbq.array_agg(df_manuals_all[["ref_x", "uri_x"]].groupby("uri_x"))["ref_x"] + .str[0] + .to_frame() + ) + df_manuals_agg["chunks"] = bbq.array_agg( + df_manuals_all[["ref_y", "uri_x"]].groupby("uri_x") + )["ref_y"] + # [END bigquery_dataframes_create_manual_to_chunks] + + # [START bigquery_dataframes_show_manual_to_chunks] + df_manuals_agg + # [END bigquery_dataframes_show_manual_to_chunks] + + # [START bigquery_dataframes_generate_pages_summary] + df_manuals_agg["chunks_url"] = bbq.array_agg( + bbq.obj.get_access_url(df_manuals_agg.explode("chunks")["chunks"], "R").groupby( + "uri_x" + ) + ) + df_manuals_agg[ + "prompt0" + ] = "Can you provide a page by page summary for the first 3 pages of the attached manual? Only write one line for each page. The pages are provided in serial order" + df_manuals_agg["prompt"] = bbq.struct(df_manuals_agg[["prompt0", "chunks_url"]]) + + result = bbq.ai.generate_text(gemini_model, df_manuals_agg["prompt"])["result"] + result + # [END bigquery_dataframes_generate_pages_summary] + + # [START bigquery_dataframes_generate_each_page_summary] + result = bbq.ai.generate_table( + gemini_model, + df_manuals_agg["prompt"], + output_schema={ + "page1_summary": "STRING", + "page2_summary": "STRING", + "page3_summary": "STRING", + }, + )[["page1_summary", "page2_summary", "page3_summary"]] + result + # [END bigquery_dataframes_generate_each_page_summary] diff --git a/samples/snippets/noxfile.py b/samples/snippets/noxfile.py new file mode 100644 index 00000000000..63e742993f9 --- /dev/null +++ b/samples/snippets/noxfile.py @@ -0,0 +1,291 @@ +# Copyright 2019 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import print_function + +import glob +import os +from pathlib import Path +import sys +from typing import Callable, Dict, Optional + +import nox + +# WARNING - WARNING - WARNING - WARNING - WARNING +# WARNING - WARNING - WARNING - WARNING - WARNING +# DO NOT EDIT THIS FILE EVER! +# WARNING - WARNING - WARNING - WARNING - WARNING +# WARNING - WARNING - WARNING - WARNING - WARNING + +BLACK_VERSION = "black==22.3.0" +ISORT_VERSION = "isort==5.10.1" + +# Copy `noxfile_config.py` to your directory and modify it instead. + +# `TEST_CONFIG` dict is a configuration hook that allows users to +# modify the test configurations. The values here should be in sync +# with `noxfile_config.py`. Users will copy `noxfile_config.py` into +# their directory and modify it. + +TEST_CONFIG = { + # You can opt out from the test for specific Python versions. + "ignored_versions": [], + # Old samples are opted out of enforcing Python type hints + # All new samples should feature them + "enforce_type_hints": False, + # An envvar key for determining the project id to use. Change it + # to 'BUILD_SPECIFIC_GCLOUD_PROJECT' if you want to opt in using a + # build specific Cloud project. You can also use your own string + # to use your own Cloud project. + "gcloud_project_env": "GOOGLE_CLOUD_PROJECT", + # 'gcloud_project_env': 'BUILD_SPECIFIC_GCLOUD_PROJECT', + # If you need to use a specific version of pip, + # change pip_version_override to the string representation + # of the version number, for example, "20.2.4" + "pip_version_override": None, + # A dictionary you want to inject into your test. Don't put any + # secrets here. These values will override predefined values. + "envs": {}, +} + + +try: + # Ensure we can import noxfile_config in the project's directory. + sys.path.append(".") + from noxfile_config import TEST_CONFIG_OVERRIDE +except ImportError as e: + print("No user noxfile_config found: detail: {}".format(e)) + TEST_CONFIG_OVERRIDE = {} + +# Update the TEST_CONFIG with the user supplied values. +TEST_CONFIG.update(TEST_CONFIG_OVERRIDE) + + +def get_pytest_env_vars() -> Dict[str, str]: + """Returns a dict for pytest invocation.""" + ret = {} + + # Override the GCLOUD_PROJECT and the alias. + env_key = TEST_CONFIG["gcloud_project_env"] + # This should error out if not set. + ret["GOOGLE_CLOUD_PROJECT"] = os.environ[env_key] + + # Apply user supplied envs. + ret.update(TEST_CONFIG["envs"]) + return ret + + +# All versions used to test samples. +ALL_VERSIONS = ["3.10", "3.11", "3.12", "3.13", "3.14"] + +# Any default versions that should be ignored. +IGNORED_VERSIONS = TEST_CONFIG["ignored_versions"] + +TESTED_VERSIONS = sorted([v for v in ALL_VERSIONS if v not in IGNORED_VERSIONS]) + +INSTALL_LIBRARY_FROM_SOURCE = os.environ.get("INSTALL_LIBRARY_FROM_SOURCE", False) in ( + "True", + "true", +) + +# Error if a python version is missing +nox.options.error_on_missing_interpreters = True + +# +# Style Checks +# + + +# Linting with flake8. +# +# We ignore the following rules: +# E203: whitespace before ‘:’ +# E266: too many leading ‘#’ for block comment +# E501: line too long +# I202: Additional newline in a section of imports +# +# We also need to specify the rules which are ignored by default: +# ['E226', 'W504', 'E126', 'E123', 'W503', 'E24', 'E704', 'E121'] +FLAKE8_COMMON_ARGS = [ + "--show-source", + "--builtin=gettext", + "--max-complexity=20", + "--exclude=.nox,.cache,env,lib,generated_pb2,*_pb2.py,*_pb2_grpc.py", + "--ignore=E121,E123,E126,E203,E226,E24,E266,E501,E704,W503,W504,I202", + "--max-line-length=88", +] + + +@nox.session +def lint(session: nox.sessions.Session) -> None: + if not TEST_CONFIG["enforce_type_hints"]: + session.install("flake8") + else: + session.install("flake8", "flake8-annotations") + + args = FLAKE8_COMMON_ARGS + [ + ".", + ] + session.run("flake8", *args) + + +# +# Black +# + + +@nox.session +def blacken(session: nox.sessions.Session) -> None: + """Run black. Format code to uniform standard.""" + session.install(BLACK_VERSION) + python_files = [path for path in os.listdir(".") if path.endswith(".py")] + + session.run("black", *python_files) + + +# +# format = isort + black +# + + +@nox.session +def format(session: nox.sessions.Session) -> None: + """ + Run isort to sort imports. Then run black + to format code to uniform standard. + """ + session.install(BLACK_VERSION, ISORT_VERSION) + python_files = [path for path in os.listdir(".") if path.endswith(".py")] + + # Use the --fss option to sort imports using strict alphabetical order. + # See https://pycqa.github.io/isort/docs/configuration/options.html#force-sort-within-sections + session.run("isort", "--fss", *python_files) + session.run("black", *python_files) + + +# +# Sample Tests +# + + +PYTEST_COMMON_ARGS = ["--junitxml=sponge_log.xml"] + + +def _session_tests( + session: nox.sessions.Session, post_install: Callable = None +) -> None: + # check for presence of tests + test_list = glob.glob("**/*_test.py", recursive=True) + glob.glob( + "**/test_*.py", recursive=True + ) + test_list.extend(glob.glob("**/tests", recursive=True)) + + if len(test_list) == 0: + print("No tests found, skipping directory.") + return + + if TEST_CONFIG["pip_version_override"]: + pip_version = TEST_CONFIG["pip_version_override"] + session.install(f"pip=={pip_version}") + """Runs py.test for a particular project.""" + concurrent_args = [] + if os.path.exists("requirements.txt"): + if os.path.exists("constraints.txt"): + session.install("-r", "requirements.txt", "-c", "constraints.txt") + else: + session.install("-r", "requirements.txt") + with open("requirements.txt") as rfile: + packages = rfile.read() + + if os.path.exists("requirements-test.txt"): + if os.path.exists("constraints-test.txt"): + session.install("-r", "requirements-test.txt", "-c", "constraints-test.txt") + else: + session.install("-r", "requirements-test.txt") + with open("requirements-test.txt") as rtfile: + packages += rtfile.read() + + if INSTALL_LIBRARY_FROM_SOURCE: + session.install("-e", _get_repo_root()) + + if post_install: + post_install(session) + + if "pytest-parallel" in packages: + concurrent_args.extend(["--workers", "auto", "--tests-per-worker", "auto"]) + elif "pytest-xdist" in packages: + concurrent_args.extend(["-n", "auto"]) + + session.run( + "pytest", + *(PYTEST_COMMON_ARGS + session.posargs + concurrent_args), + # Pytest will return 5 when no tests are collected. This can happen + # on travis where slow and flaky tests are excluded. + # See http://doc.pytest.org/en/latest/_modules/_pytest/main.html + success_codes=[0, 5], + env=get_pytest_env_vars(), + ) + + +@nox.session(python=ALL_VERSIONS) +def py(session: nox.sessions.Session) -> None: + """Runs py.test for a sample using the specified version of Python.""" + if session.python in TESTED_VERSIONS: + _session_tests(session) + else: + session.skip( + "SKIPPED: {} tests are disabled for this sample.".format(session.python) + ) + + +# +# Readmegen +# + + +def _get_repo_root() -> Optional[str]: + """Returns the root folder of the project.""" + # Get root of this repository. Assume we don't have directories nested deeper than 10 items. + p = Path(os.getcwd()) + for i in range(10): + if p is None: + break + if Path(p / ".git").exists(): + return str(p) + # .git is not available in repos cloned via Cloud Build + # setup.py is always in the library's root, so use that instead + # https://github.com/googleapis/synthtool/issues/792 + if Path(p / "setup.py").exists(): + return str(p) + p = p.parent + raise Exception("Unable to detect repository root.") + + +GENERATED_READMES = sorted([x for x in Path(".").rglob("*.rst.in")]) + + +@nox.session +@nox.parametrize("path", GENERATED_READMES) +def readmegen(session: nox.sessions.Session, path: str) -> None: + """(Re-)generates the readme for a sample.""" + session.install("jinja2", "pyyaml") + dir_ = os.path.dirname(path) + + if os.path.exists(os.path.join(dir_, "requirements.txt")): + session.install("-r", os.path.join(dir_, "requirements.txt")) + + in_file = os.path.join(dir_, "README.rst.in") + session.run( + "python", _get_repo_root() + "/scripts/readme-gen/readme_gen.py", in_file + ) diff --git a/samples/snippets/noxfile_config.py b/samples/snippets/noxfile_config.py new file mode 100644 index 00000000000..211d6974b9e --- /dev/null +++ b/samples/snippets/noxfile_config.py @@ -0,0 +1,42 @@ +# Copyright 2020 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Default TEST_CONFIG_OVERRIDE for python repos. + +# You can copy this file into your directory, then it will be inported from +# the noxfile.py. + +# The source of truth: +# https://github.com/GoogleCloudPlatform/python-docs-samples/blob/master/noxfile_config.py + +TEST_CONFIG_OVERRIDE = { + # You can opt out from the test for specific Python versions. + "ignored_versions": ["2.7", "3.7", "3.8"], + # Old samples are opted out of enforcing Python type hints + # All new samples should feature them + "enforce_type_hints": True, + # An envvar key for determining the project id to use. Change it + # to 'BUILD_SPECIFIC_GCLOUD_PROJECT' if you want to opt in using a + # build specific Cloud project. You can also use your own string + # to use your own Cloud project. + "gcloud_project_env": "GOOGLE_CLOUD_PROJECT", + # "gcloud_project_env": "BUILD_SPECIFIC_GCLOUD_PROJECT", + # If you need to use a specific version of pip, + # change pip_version_override to the string representation + # of the version number, for example, "20.2.4" + "pip_version_override": None, + # A dictionary you want to inject into your test. Don't put any + # secrets here. These values will override predefined values. + "envs": {}, +} diff --git a/samples/snippets/ordering_mode_partial_test.py b/samples/snippets/ordering_mode_partial_test.py new file mode 100644 index 00000000000..38b4bf15278 --- /dev/null +++ b/samples/snippets/ordering_mode_partial_test.py @@ -0,0 +1,50 @@ +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +def test_quickstart() -> None: + import bigframes.pandas + + try: + # We need a fresh session since we're modifying connection options. + bigframes.pandas.close_session() + + # [START bigquery_bigframes_ordering_mode_partial] + import bigframes.pandas as bpd + + bpd.options.bigquery.ordering_mode = "partial" + # [END bigquery_bigframes_ordering_mode_partial] + + # [START bigquery_bigframes_ordering_mode_partial_ambiguous_window_warning] + import warnings + + import bigframes.exceptions + + warnings.simplefilter( + "ignore", category=bigframes.exceptions.AmbiguousWindowWarning + ) + # [END bigquery_bigframes_ordering_mode_partial_ambiguous_window_warning] + + df = bpd.DataFrame({"column": [1, 2, 1, 3, 1, 2, 3]}) + + # [START bigquery_bigframes_ordering_mode_partial_drop_duplicates] + # Avoid order dependency by using groupby instead of drop_duplicates. + unique_col = df.groupby(["column"], as_index=False).size().drop(columns="size") + # [END bigquery_bigframes_ordering_mode_partial_drop_duplicates] + + assert len(unique_col) == 3 + finally: + # Don't leak partial ordering mode to other code samples. + bigframes.pandas.close_session() + bpd.options.bigquery.ordering_mode = "strict" diff --git a/samples/snippets/pandas_methods_test.py b/samples/snippets/pandas_methods_test.py new file mode 100644 index 00000000000..0f128f9e6ac --- /dev/null +++ b/samples/snippets/pandas_methods_test.py @@ -0,0 +1,41 @@ +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +def test_bigquery_dataframes_pandas_methods() -> None: + # [START bigquery_dataframes_pandas_methods] + import bigframes.pandas as bpd + + # Load data from BigQuery + query_or_table = "bigquery-public-data.ml_datasets.penguins" + bq_df = bpd.read_gbq(query_or_table) + + # Inspect one of the columns (or series) of the DataFrame: + bq_df["body_mass_g"] + + # Compute the mean of this series: + average_body_mass = bq_df["body_mass_g"].mean() + print(f"average_body_mass: {average_body_mass}") + + # Find the heaviest species using the groupby operation to calculate the + # mean body_mass_g: + ( + bq_df["body_mass_g"] + .groupby(by=bq_df["species"]) + .mean() + .sort_values(ascending=False) + .head(10) + ) + # [END bigquery_dataframes_pandas_methods] + assert average_body_mass is not None diff --git a/samples/snippets/performance_optimizations_test.py b/samples/snippets/performance_optimizations_test.py new file mode 100644 index 00000000000..43e14e31cc0 --- /dev/null +++ b/samples/snippets/performance_optimizations_test.py @@ -0,0 +1,52 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +def test_performance_optimizations() -> None: + # [START bigquery_bigframes_use_peek_to_preview_data] + import bigframes.pandas as bpd + + # Read the "Penguins" table into a dataframe + df = bpd.read_gbq("bigquery-public-data.ml_datasets.penguins") + + # Preview 3 random rows + df.peek(3) + # [END bigquery_bigframes_use_peek_to_preview_data] + assert df.peek(3) is not None + + import bigframes.pandas as bpd + + users = bpd.DataFrame({"user_name": ["John"]}) + groups = bpd.DataFrame({"group_id": ["group_1"]}) + transactions = bpd.DataFrame({"amount": [3], "completed": [True]}) + + # [START bigquery_bigframes_use_cache_after_expensive_operations] + # Assume you have 3 large dataframes "users", "group" and "transactions" + + # Expensive join operations + final_df = users.join(groups).join(transactions) + final_df.cache() + # Subsequent derived results will reuse the cached join + print(final_df.peek()) + print(len(final_df[final_df["completed"]])) + print(final_df.groupby("group_id")["amount"].mean().peek(30)) + # [END bigquery_bigframes_use_cache_after_expensive_operations] + assert final_df is not None + + # [START bigquery_bigframes_enable_deferred_repr_for_debugging] + import bigframes.pandas as bpd + + bpd.options.display.repr_mode = "deferred" + # [END bigquery_bigframes_enable_deferred_repr_for_debugging] + assert bpd.options.display.repr_mode == "deferred" diff --git a/samples/snippets/quickstart.py b/samples/snippets/quickstart.py new file mode 100644 index 00000000000..08662c1ea7e --- /dev/null +++ b/samples/snippets/quickstart.py @@ -0,0 +1,81 @@ +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +def run_quickstart(project_id: str) -> None: + your_gcp_project_id = project_id + + # [START bigquery_bigframes_quickstart_create_dataframe] + import bigframes.pandas as bpd + + # Set BigQuery DataFrames options + # Note: The project option is not required in all environments. + # On BigQuery Studio, the project ID is automatically detected. + bpd.options.bigquery.project = your_gcp_project_id + + # Use "partial" ordering mode to generate more efficient queries, but the + # order of the rows in DataFrames may not be deterministic if you have not + # explictly sorted it. Some operations that depend on the order, such as + # head() will not function until you explictly order the DataFrame. Set the + # ordering mode to "strict" (default) for more pandas compatibility. + bpd.options.bigquery.ordering_mode = "partial" + + # Create a DataFrame from a BigQuery table + query_or_table = "bigquery-public-data.ml_datasets.penguins" + df = bpd.read_gbq(query_or_table) + + # Efficiently preview the results using the .peek() method. + df.peek() + # [END bigquery_bigframes_quickstart_create_dataframe] + + # [START bigquery_bigframes_quickstart_calculate_print] + # Use the DataFrame just as you would a pandas DataFrame, but calculations + # happen in the BigQuery query engine instead of the local system. + average_body_mass = df["body_mass_g"].mean() + print(f"average_body_mass: {average_body_mass}") + # [END bigquery_bigframes_quickstart_calculate_print] + + # [START bigquery_bigframes_quickstart_eval_metrics] + # Create the Linear Regression model + from bigframes.ml.linear_model import LinearRegression + + # Filter down to the data we want to analyze + adelie_data = df[df.species == "Adelie Penguin (Pygoscelis adeliae)"] + + # Drop the columns we don't care about + adelie_data = adelie_data.drop(columns=["species"]) + + # Drop rows with nulls to get our training data + training_data = adelie_data.dropna() + + # Pick feature columns and label column + X = training_data[ + [ + "island", + "culmen_length_mm", + "culmen_depth_mm", + "flipper_length_mm", + "sex", + ] + ] + y = training_data[["body_mass_g"]] + + model = LinearRegression(fit_intercept=False) + model.fit(X, y) + model.score(X, y) + # [END bigquery_bigframes_quickstart_eval_metrics] + + # close session and reset option so not to affect other tests + bpd.close_session() + bpd.options.reset() diff --git a/samples/snippets/quickstart_test.py b/samples/snippets/quickstart_test.py new file mode 100644 index 00000000000..a650f8365dc --- /dev/null +++ b/samples/snippets/quickstart_test.py @@ -0,0 +1,36 @@ +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + +import bigframes.pandas + +from . import quickstart + + +def test_quickstart( + capsys: pytest.CaptureFixture[str], +) -> None: + # We need a fresh session since we're modifying connection options. + bigframes.pandas.close_session() + + # Determine project id, in this case prefer the one set in the environment + # variable GOOGLE_CLOUD_PROJECT (if any) + import os + + your_project_id = os.getenv("GOOGLE_CLOUD_PROJECT", "bigframes-dev") + + quickstart.run_quickstart(your_project_id) + out, _ = capsys.readouterr() + assert "average_body_mass:" in out diff --git a/samples/snippets/regression_model_test.py b/samples/snippets/regression_model_test.py new file mode 100644 index 00000000000..43cdabacb46 --- /dev/null +++ b/samples/snippets/regression_model_test.py @@ -0,0 +1,57 @@ +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +def test_regression_model() -> None: + # [START bigquery_dataframes_regression_model] + from bigframes.ml.linear_model import LinearRegression + import bigframes.pandas as bpd + + # Load data from BigQuery + query_or_table = "bigquery-public-data.ml_datasets.penguins" + bq_df = bpd.read_gbq(query_or_table) + + # Filter down to the data to the Adelie Penguin species + adelie_data = bq_df[bq_df.species == "Adelie Penguin (Pygoscelis adeliae)"] + + # Drop the species column + adelie_data = adelie_data.drop(columns=["species"]) + + # Drop rows with nulls to get training data + training_data = adelie_data.dropna() + + # Specify your feature (or input) columns and the label (or output) column: + feature_columns = training_data[ + ["island", "culmen_length_mm", "culmen_depth_mm", "flipper_length_mm", "sex"] + ] + label_columns = training_data[["body_mass_g"]] + + test_data = adelie_data[adelie_data.body_mass_g.isnull()] + + # Create the linear model + model = LinearRegression() + model.fit(feature_columns, label_columns) + + # Score the model + score = model.score(feature_columns, label_columns) + + # Predict using the model + result = model.predict(test_data) + # [END bigquery_dataframes_regression_model] + assert test_data is not None + assert feature_columns is not None + assert label_columns is not None + assert model is not None + assert score is not None + assert result is not None diff --git a/samples/snippets/remote_function.py b/samples/snippets/remote_function.py new file mode 100644 index 00000000000..4c5b365007a --- /dev/null +++ b/samples/snippets/remote_function.py @@ -0,0 +1,159 @@ +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +def run_remote_function_and_read_gbq_function(project_id: str) -> None: + your_gcp_project_id = project_id + + # [START bigquery_dataframes_remote_function] + import bigframes.pandas as bpd + + # Set BigQuery DataFrames options + bpd.options.bigquery.project = your_gcp_project_id + bpd.options.bigquery.location = "US" + + # BigQuery DataFrames gives you the ability to turn your custom scalar + # functions into a BigQuery remote function. It requires the GCP project to + # be set up appropriately and the user having sufficient privileges to use + # them. One can find more details about the usage and the requirements via + # `help` command. + help(bpd.remote_function) + + # Read a table and inspect the column of interest. + df = bpd.read_gbq("bigquery-public-data.ml_datasets.penguins") + df["body_mass_g"].head(10) + + # Define a custom function, and specify the intent to turn it into a remote + # function. It requires a BigQuery connection. If the connection is not + # already created, BigQuery DataFrames will attempt to create one assuming + # the necessary APIs and IAM permissions are setup in the project. In our + # examples we will be letting the default connection `bigframes-default-connection` + # be used. We will also set `reuse=False` to make sure we don't + # step over someone else creating remote function in the same project from + # the exact same source code at the same time. Let's try a `pandas`-like use + # case in which we want to apply a user defined scalar function to every + # value in a `Series`, more specifically bucketize the `body_mass_g` value + # of the penguins, which is a real number, into a category, which is a + # string. + @bpd.remote_function( + reuse=False, + cloud_function_service_account="default", + ) + def get_bucket(num: float) -> str: + if not num: + return "NA" + boundary = 4000 + return "at_or_above_4000" if num >= boundary else "below_4000" + + # Then we can apply the remote function on the `Series` of interest via + # `apply` API and store the result in a new column in the DataFrame. + df = df.assign(body_mass_bucket=df["body_mass_g"].apply(get_bucket)) + + # This will add a new column `body_mass_bucket` in the DataFrame. You can + # preview the original value and the bucketized value side by side. + df[["body_mass_g", "body_mass_bucket"]].head(10) + + # The above operation was possible by doing all the computation on the + # cloud. For that, there is a google cloud function deployed by serializing + # the user code, and a BigQuery remote function created to call the cloud + # function via the latter's http endpoint on the data in the DataFrame. + + # The BigQuery remote function created to support the BigQuery DataFrames + # remote function can be located via a property `bigframes_remote_function` + # set in the remote function object. + print(f"Created BQ remote function: {get_bucket.bigframes_remote_function}") + + # The cloud function can be located via another property + # `bigframes_cloud_function` set in the remote function object. + print(f"Created cloud function: {get_bucket.bigframes_cloud_function}") + + # Warning: The deployed cloud function may be visible to other users with + # sufficient privilege in the project, so the user should be careful about + # having any sensitive data in the code that will be deployed as a remote + # function. + + # Let's continue trying other potential use cases of remote functions. Let's + # say we consider the `species`, `island` and `sex` of the penguins + # sensitive information and want to redact that by replacing with their hash + # code instead. Let's define another scalar custom function and decorate it + # as a remote function. The custom function in this example has external + # package dependency, which can be specified via `packages` parameter. + @bpd.remote_function( + reuse=False, + packages=["cryptography"], + cloud_function_service_account="default", + ) + def get_hash(input: str) -> str: + from cryptography.fernet import Fernet + + # handle missing value + if input is None: + input = "" + + key = Fernet.generate_key() + f = Fernet(key) + return f.encrypt(input.encode()).decode() + + # We can use this remote function in another `pandas`-like API `map` that + # can be applied on a DataFrame + df_redacted = df[["species", "island", "sex"]].map(get_hash) + df_redacted.head(10) + + # [END bigquery_dataframes_remote_function] + + existing_get_bucket_bq_udf = get_bucket.bigframes_remote_function + + # [START bigquery_dataframes_read_gbq_function] + + # If you have already defined a custom function in BigQuery, either via the + # BigQuery Google Cloud Console or with the `remote_function` decorator, + # or otherwise, you may use it with BigQuery DataFrames with the + # `read_gbq_function` method. More details are available via the `help` + # command. + import bigframes.pandas as pd + + help(pd.read_gbq_function) + + # Here is an example of using `read_gbq_function` to load an existing + # BigQuery function. + df = pd.read_gbq("bigquery-public-data.ml_datasets.penguins") + get_bucket_function = pd.read_gbq_function(existing_get_bucket_bq_udf) + + df = df.assign(body_mass_bucket=df["body_mass_g"].apply(get_bucket_function)) + df.head(10) + + # It should be noted that if a function is created using the + # `remote_function` decorator, its created BQ remote function is accessible + # immediately afterward via the function's `bigframes_remote_function` + # attribute. The same string can be passed to `read_gbq_function` later in + # another context. + + # [END bigquery_dataframes_read_gbq_function] + + # Clean up cloud artifacts + session = bpd.get_global_session() + for function in (get_bucket, get_hash): + try: + session.bqclient.delete_routine(function.bigframes_remote_function) + except Exception: + # Ignore exception during clean-up + pass + + try: + session.cloudfunctionsclient.delete_function( + name=function.bigframes_cloud_function + ) + except Exception: + # Ignore exception during clean-up + pass diff --git a/samples/snippets/remote_function_test.py b/samples/snippets/remote_function_test.py new file mode 100644 index 00000000000..24bc7e854e2 --- /dev/null +++ b/samples/snippets/remote_function_test.py @@ -0,0 +1,68 @@ +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import google.api_core.exceptions +import google.cloud.bigquery_connection_v1 +import pytest + +import bigframes.pandas + +from . import remote_function + + +# TODO(tswast): Once the connections are cleaned up in the sample test project +# and https://github.com/GoogleCloudPlatform/python-docs-samples/issues/11720 +# is closed, we shouldn't need this because AFAIK we only use one BQ connection +# in this sample. +@pytest.fixture(autouse=True) +def cleanup_connections() -> None: + client = google.cloud.bigquery_connection_v1.ConnectionServiceClient() + + for conn in client.list_connections( + parent="projects/python-docs-samples-tests/locations/us" + ): + try: + int(conn.name.split("/")[-1].split("-")[0], base=16) + except ValueError: + print(f"Couldn't parse {conn.name}") + continue + + try: + print(f"removing {conn.name}") + client.delete_connection( + google.cloud.bigquery_connection_v1.DeleteConnectionRequest( + {"name": conn.name}, + ) + ) + except google.api_core.exceptions.GoogleAPIError: + # We did as much clean up as we can. + break + + +def test_remote_function_and_read_gbq_function( + capsys: pytest.CaptureFixture[str], +) -> None: + # We need a fresh session since we're modifying connection options. + bigframes.pandas.close_session() + + # Determine project id, in this case prefer the one set in the environment + # variable GOOGLE_CLOUD_PROJECT (if any) + import os + + your_project_id = os.getenv("GOOGLE_CLOUD_PROJECT", "bigframes-dev") + + remote_function.run_remote_function_and_read_gbq_function(your_project_id) + out, _ = capsys.readouterr() + assert "Created BQ remote function:" in out + assert "Created cloud function:" in out diff --git a/samples/snippets/requirements-test.txt b/samples/snippets/requirements-test.txt new file mode 100644 index 00000000000..cbac5e3f126 --- /dev/null +++ b/samples/snippets/requirements-test.txt @@ -0,0 +1,3 @@ +# samples/snippets should be runnable with no "extras" +google-cloud-testutils==1.4.0 +pytest==8.3.2 diff --git a/samples/snippets/requirements.txt b/samples/snippets/requirements.txt new file mode 100644 index 00000000000..9b5da5182e2 --- /dev/null +++ b/samples/snippets/requirements.txt @@ -0,0 +1,2 @@ +# samples/snippets should be runnable with no "extras" +bigframes==1.11.1 diff --git a/samples/snippets/sessions_and_io_test.py b/samples/snippets/sessions_and_io_test.py new file mode 100644 index 00000000000..06f0c4ab3c3 --- /dev/null +++ b/samples/snippets/sessions_and_io_test.py @@ -0,0 +1,180 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +def test_sessions_and_io(project_id: str, dataset_id: str, gcs_bucket: str) -> None: + YOUR_PROJECT_ID = project_id + YOUR_DATASET_ID = dataset_id + YOUR_LOCATION = "us" + YOUR_BUCKET = gcs_bucket + + # [START bigquery_dataframes_create_and_use_session_instance] + import bigframes + import bigframes.pandas as bpd + + # Create session object + context = bigframes.BigQueryOptions( + project=YOUR_PROJECT_ID, + location=YOUR_LOCATION, + ) + session = bigframes.Session(context) + + # Load a BigQuery table into a dataframe + df1 = session.read_gbq("bigquery-public-data.ml_datasets.penguins") + + # Create a dataframe with local data: + df2 = bpd.DataFrame({"my_col": [1, 2, 3]}, session=session) + # [END bigquery_dataframes_create_and_use_session_instance] + assert df1 is not None + assert df2 is not None + + # [START bigquery_dataframes_combine_data_from_multiple_sessions_raise_error] + import bigframes + import bigframes.pandas as bpd + + context = bigframes.BigQueryOptions(location=YOUR_LOCATION, project=YOUR_PROJECT_ID) + + session1 = bigframes.Session(context) + session2 = bigframes.Session(context) + + series1 = bpd.Series([1, 2, 3, 4, 5], session=session1) + series2 = bpd.Series([1, 2, 3, 4, 5], session=session2) + + try: + series1 + series2 + except ValueError as e: + print(e) # Error message: Cannot use combine sources from multiple sessions + # [END bigquery_dataframes_combine_data_from_multiple_sessions_raise_error] + + # [START bigquery_dataframes_set_options_for_global_session] + import bigframes.pandas as bpd + + # Set project ID for the global session + bpd.options.bigquery.project = YOUR_PROJECT_ID + # Update the global default session location + bpd.options.bigquery.location = YOUR_LOCATION + # [END bigquery_dataframes_set_options_for_global_session] + + # [START bigquery_dataframes_global_session_is_the_default_session] + # The following two statements are essentially the same + df = bpd.read_gbq("bigquery-public-data.ml_datasets.penguins") + df = bpd.get_global_session().read_gbq("bigquery-public-data.ml_datasets.penguins") + # [END bigquery_dataframes_global_session_is_the_default_session] + assert df is not None + + # [START bigquery_dataframes_create_dataframe_from_py_and_np] + import numpy as np + + import bigframes.pandas as bpd + + s = bpd.Series([1, 2, 3]) + + # Create a dataframe with Python dict + df = bpd.DataFrame( + { + "col_1": [1, 2, 3], + "col_2": [4, 5, 6], + } + ) + + # Create a series with Numpy + s = bpd.Series(np.arange(10)) + # [END bigquery_dataframes_create_dataframe_from_py_and_np] + assert s is not None + + # [START bigquery_dataframes_create_dataframe_from_pandas] + import numpy as np + import pandas as pd + + import bigframes.pandas as bpd + + pd_df = pd.DataFrame(np.random.randn(4, 2)) + + # Convert Pandas dataframe to BigQuery DataFrame with read_pandas() + df_1 = bpd.read_pandas(pd_df) + # Convert Pandas dataframe to BigQuery DataFrame with the dataframe constructor + df_2 = bpd.DataFrame(pd_df) + # [END bigquery_dataframes_create_dataframe_from_pandas] + assert df_1 is not None + assert df_2 is not None + + # [START bigquery_dataframes_convert_bq_dataframe_to_pandas] + import bigframes.pandas as bpd + + bf_df = bpd.DataFrame({"my_col": [1, 2, 3]}) + # Returns a Pandas Dataframe + bf_df.to_pandas() + + bf_s = bpd.Series([1, 2, 3]) + # Returns a Pandas Series + bf_s.to_pandas() + # [END bigquery_dataframes_convert_bq_dataframe_to_pandas] + assert bf_s.to_pandas() is not None + + # [START bigquery_dataframes_to_pandas_dry_run] + import bigframes.pandas as bpd + + df = bpd.read_gbq("bigquery-public-data.ml_datasets.penguins") + + # Returns a Pandas series with dry run stats + df.to_pandas(dry_run=True) + # [END bigquery_dataframes_to_pandas_dry_run] + assert df.to_pandas(dry_run=True) is not None + + # [START bigquery_dataframes_read_data_from_csv] + import bigframes.pandas as bpd + + # Read a CSV file from GCS + df = bpd.read_csv("gs://cloud-samples-data/bigquery/us-states/us-states.csv") + # [END bigquery_dataframes_read_data_from_csv] + assert df is not None + + # [START bigquery_dataframes_write_data_to_csv] + import bigframes.pandas as bpd + + df = bpd.DataFrame({"my_col": [1, 2, 3]}) + # Write a dataframe to a CSV file in GCS + df.to_csv(f"gs://{YOUR_BUCKET}/myfile*.csv") + # [END bigquery_dataframes_write_data_to_csv] + assert df is not None + + # [START bigquery_dataframes_read_data_from_bigquery_table] + import bigframes.pandas as bpd + + df = bpd.read_gbq("bigquery-public-data.ml_datasets.penguins") + # [END bigquery_dataframes_read_data_from_bigquery_table] + assert df is not None + + # [START bigquery_dataframes_read_from_sql_query] + import bigframes.pandas as bpd + + sql = """ + SELECT species, island, body_mass_g + FROM bigquery-public-data.ml_datasets.penguins + WHERE sex = 'MALE' + """ + + df = bpd.read_gbq(sql) + # [END bigquery_dataframes_read_from_sql_query] + assert df is not None + + YOUR_TABLE_NAME = "snippets-session-and-io-test" + + # [START bigquery_dataframes_dataframe_to_bigquery_table] + import bigframes.pandas as bpd + + df = bpd.DataFrame({"my_col": [1, 2, 3]}) + + df.to_gbq(f"{YOUR_PROJECT_ID}.{YOUR_DATASET_ID}.{YOUR_TABLE_NAME}") + # [END bigquery_dataframes_dataframe_to_bigquery_table] diff --git a/samples/snippets/set_options_test.py b/samples/snippets/set_options_test.py new file mode 100644 index 00000000000..6007dcbb383 --- /dev/null +++ b/samples/snippets/set_options_test.py @@ -0,0 +1,45 @@ +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +def test_bigquery_dataframes_set_options() -> None: + # Close the session before resetting the options + import bigframes.pandas as bpd + + bpd.close_session() + + try: + # [START bigquery_dataframes_set_options] + import bigframes.pandas as bpd + + PROJECT_ID = "bigframes-dev" # @param {type:"string"} + REGION = "US" # @param {type:"string"} + + # Set BigQuery DataFrames options + # Note: The project option is not required in all environments. + # On BigQuery Studio, the project ID is automatically detected. + bpd.options.bigquery.project = PROJECT_ID + + # Note: The location option is not required. + # It defaults to the location of the first table or query + # passed to read_gbq(). For APIs where a location can't be + # auto-detected, the location defaults to the "US" location. + bpd.options.bigquery.location = REGION + + # [END bigquery_dataframes_set_options] + assert bpd.options.bigquery.project == PROJECT_ID + assert bpd.options.bigquery.location == REGION + finally: + bpd.close_session() + bpd.options.reset() diff --git a/samples/snippets/st_regionstats_test.py b/samples/snippets/st_regionstats_test.py new file mode 100644 index 00000000000..f0f4963a824 --- /dev/null +++ b/samples/snippets/st_regionstats_test.py @@ -0,0 +1,80 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Code sample for https://docs.cloud.google.com/bigquery/docs/raster-data#analytics-hub-source""" + + +def test_st_regionstats() -> None: + project_id = "bigframes-dev" + + # [START bigquery_dataframes_st_regionstats] + import datetime + from typing import cast + + import bigframes.bigquery as bbq + import bigframes.pandas as bpd + + # TODO: Set the project_id to your Google Cloud project ID. + # project_id = "your-project-id" + bpd.options.bigquery.project = project_id + + # TODO: Set the dataset_id to the ID of the dataset that contains the + # `climate` table. This is likely a linked dataset to Earth Engine. + # See: https://cloud.google.com/bigquery/docs/link-earth-engine + linked_dataset = "era5_land_daily_aggregated" + + # For the best efficiency, use partial ordering mode. + bpd.options.bigquery.ordering_mode = "partial" + + # Load the table of country boundaries. + countries = bpd.read_gbq("bigquery-public-data.overture_maps.division_area") + + # Filter to just the countries. + countries = countries[countries["subtype"] == "country"].copy() + countries["name"] = countries["names"].struct.field("primary") + countries["simplified_geometry"] = bbq.st_simplify( + countries["geometry"], + tolerance_meters=10_000, + ) + + # Get the reference to the temperature data from a linked dataset. + # Note: This sample assumes you have a linked dataset to Earth Engine. + image_href = ( + bpd.read_gbq(f"{project_id}.{linked_dataset}.climate") + .set_index("start_datetime") + .loc[[datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc)], :] + ) + raster_id = image_href["assets"].struct.field("image").struct.field("href") + raster_id = raster_id.item() + stats = bbq.st_regionstats( + countries["simplified_geometry"], + raster_id=cast(str, raster_id), + band="temperature_2m", + ) + + # Extract the mean and convert from Kelvin to Celsius. + countries["mean_temperature"] = stats.struct.field("mean") - 273.15 + + # Sort by the mean temperature to find the warmest countries. + result = countries[["name", "mean_temperature"]].sort_values( + "mean_temperature", ascending=False + ) + print(result.head(10)) + # [END bigquery_dataframes_st_regionstats] + + assert len(result) > 0 + + +if __name__ == "__main__": + test_st_regionstats() diff --git a/samples/snippets/type_system_test.py b/samples/snippets/type_system_test.py new file mode 100644 index 00000000000..88b9e747424 --- /dev/null +++ b/samples/snippets/type_system_test.py @@ -0,0 +1,235 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pandas.testing + +from bigframes import dtypes + + +def test_type_system_examples() -> None: + # [START bigquery_dataframes_type_sytem_timestamp_local_type_conversion] + import pandas as pd + + import bigframes.pandas as bpd + + s = pd.Series([pd.Timestamp("20250101")]) + assert s.dtype == "datetime64[ns]" + assert bpd.read_pandas(s).dtype == "timestamp[us][pyarrow]" + # [END bigquery_dataframes_type_sytem_timestamp_local_type_conversion] + + # [START bigquery_dataframes_type_system_pyarrow_preference] + import datetime + + import pandas as pd + + import bigframes.pandas as bpd + + s = pd.Series([datetime.date(2025, 1, 1)]) + s + pd.Timedelta(hours=12) + # 0 2025-01-01 + # dtype: object + + bpd.read_pandas(s) + pd.Timedelta(hours=12) + # 0 2025-01-01 12:00:00 + # dtype: timestamp[us][pyarrow] + # [END bigquery_dataframes_type_system_pyarrow_preference] + pandas.testing.assert_series_equal( + s + pd.Timedelta(hours=12), pd.Series([datetime.date(2025, 1, 1)]) + ) + pandas.testing.assert_series_equal( + (bpd.read_pandas(s) + pd.Timedelta(hours=12)).to_pandas(), + pd.Series([pd.Timestamp(2025, 1, 1, 12)], dtype=dtypes.DATETIME_DTYPE), + check_index_type=False, + ) + + # [START bigquery_dataframes_type_system_load_timedelta] + import pandas as pd + + import bigframes.pandas as bpd + + s = pd.Series([pd.Timedelta("1s"), pd.Timedelta("2m")]) + bpd.read_pandas(s) + # 0 0 days 00:00:01 + # 1 0 days 00:02:00 + # dtype: duration[us][pyarrow] + # [END bigquery_dataframes_type_system_load_timedelta] + pandas.testing.assert_series_equal( + bpd.read_pandas(s).to_pandas(), + s.astype(dtypes.TIMEDELTA_DTYPE), + check_index_type=False, + ) + + # [START bigquery_dataframes_type_system_timedelta_precision] + import pandas as pd + + s = pd.Series([pd.Timedelta("999ns")]) + bpd.read_pandas(s.dt.round("us")) + # 0 0 days 00:00:00.000001 + # dtype: duration[us][pyarrow] + # [END bigquery_dataframes_type_system_timedelta_precision] + pandas.testing.assert_series_equal( + bpd.read_pandas(s.dt.round("us")).to_pandas(), + s.dt.round("us").astype(dtypes.TIMEDELTA_DTYPE), + check_index_type=False, + ) + + # [START bigquery_dataframes_type_system_cast_timedelta] + import bigframes.pandas as bpd + + bpd.to_timedelta([1, 2, 3], unit="s") + # 0 0 days 00:00:01 + # 1 0 days 00:00:02 + # 2 0 days 00:00:03 + # dtype: duration[us][pyarrow] + # [END bigquery_dataframes_type_system_cast_timedelta] + pandas.testing.assert_series_equal( + bpd.to_timedelta([1, 2, 3], unit="s").to_pandas(), + pd.Series(pd.to_timedelta([1, 2, 3], unit="s"), dtype=dtypes.TIMEDELTA_DTYPE), + check_index_type=False, + ) + + # [START bigquery_dataframes_type_system_list_accessor] + import bigframes.pandas as bpd + + s = bpd.Series([[1, 2, 3], [4, 5], [6]]) # dtype: list[pyarrow] + + # Access the first elements of each list + s.list[0] + # 0 1 + # 1 4 + # 2 6 + # dtype: Int64 + + # Get the lengths of each list + s.list.len() + # 0 3 + # 1 2 + # 2 1 + # dtype: Int64 + # [END bigquery_dataframes_type_system_list_accessor] + pandas.testing.assert_series_equal( + s.list[0].to_pandas(), + pd.Series([1, 4, 6], dtype="Int64"), + check_index_type=False, + ) + pandas.testing.assert_series_equal( + s.list.len().to_pandas(), + pd.Series([3, 2, 1], dtype="Int64"), + check_index_type=False, + ) + + # [START bigquery_dataframes_type_system_struct_accessor] + import bigframes.pandas as bpd + + structs = [ + {"id": 101, "category": "A"}, + {"id": 102, "category": "B"}, + {"id": 103, "category": "C"}, + ] + s = bpd.Series(structs) + # Get the 'id' field of each struct + s.struct.field("id") + # 0 101 + # 1 102 + # 2 103 + # Name: id, dtype: Int64 + # [END bigquery_dataframes_type_system_struct_accessor] + + # [START bigquery_dataframes_type_system_struct_accessor_shortcut] + import bigframes.pandas as bpd + + structs = [ + {"id": 101, "category": "A"}, + {"id": 102, "category": "B"}, + {"id": 103, "category": "C"}, + ] + s = bpd.Series(structs) + + # not explicitly using the "struct" property + s.id + # 0 101 + # 1 102 + # 2 103 + # Name: id, dtype: Int64 + # [END bigquery_dataframes_type_system_struct_accessor_shortcut] + pandas.testing.assert_series_equal( + s.struct.field("id").to_pandas(), + pd.Series([101, 102, 103], dtype="Int64", name="id"), + check_index_type=False, + ) + pandas.testing.assert_series_equal( + s.id.to_pandas(), + pd.Series([101, 102, 103], dtype="Int64", name="id"), + check_index_type=False, + ) + + # [START bigquery_dataframes_type_system_string_accessor] + import bigframes.pandas as bpd + + s = bpd.Series(["abc", "de", "1"]) # dtype: string[pyarrow] + + # Get the first character of each string + s.str[0] + # 0 a + # 1 d + # 2 1 + # dtype: string + + # Check whether there are only alphabetic characters in each string + s.str.isalpha() + # 0 True + # 1 True + # 2 False + # dtype: boolean + + # Cast the alphabetic characters to their upper cases for each string + s.str.upper() + # 0 ABC + # 1 DE + # 2 1 + # dtype: string + # [END bigquery_dataframes_type_system_string_accessor] + pandas.testing.assert_series_equal( + s.str[0].to_pandas(), + pd.Series(["a", "d", "1"], dtype=dtypes.STRING_DTYPE), + check_index_type=False, + ) + pandas.testing.assert_series_equal( + s.str.isalpha().to_pandas(), + pd.Series([True, True, False], dtype=dtypes.BOOL_DTYPE), + check_index_type=False, + ) + pandas.testing.assert_series_equal( + s.str.upper().to_pandas(), + pd.Series(["ABC", "DE", "1"], dtype=dtypes.STRING_DTYPE), + check_index_type=False, + ) + + # [START bigquery_dataframes_type_system_geo_accessor] + from shapely.geometry import Point + + import bigframes.pandas as bpd + + s = bpd.Series([Point(1, 0), Point(2, 1)]) # dtype: geometry + + s.geo.y + # 0 0.0 + # 1 1.0 + # dtype: Float64 + # [END bigquery_dataframes_type_system_geo_accessor] + pandas.testing.assert_series_equal( + s.geo.y.to_pandas(), + pd.Series([0.0, 1.0], dtype=dtypes.FLOAT_DTYPE), + check_index_type=False, + ) diff --git a/samples/snippets/udf.py b/samples/snippets/udf.py new file mode 100644 index 00000000000..5f7ad8a33f3 --- /dev/null +++ b/samples/snippets/udf.py @@ -0,0 +1,120 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +def run_udf_and_read_gbq_function( + project_id: str, dataset_id: str, routine_id: str +) -> None: + your_gcp_project_id = project_id + your_bq_dataset_id = dataset_id + your_bq_routine_id = routine_id + + # [START bigquery_dataframes_udf] + import bigframes.pandas as bpd + + # Set BigQuery DataFrames options + bpd.options.bigquery.project = your_gcp_project_id + bpd.options.bigquery.location = "US" + + # BigQuery DataFrames gives you the ability to turn your custom functions + # into a BigQuery Python UDF. One can find more details about the usage and + # the requirements via `help` command. + help(bpd.udf) + + # Read a table and inspect the column of interest. + df = bpd.read_gbq("bigquery-public-data.ml_datasets.penguins") + df["body_mass_g"].peek(10) + + # Define a custom function, and specify the intent to turn it into a + # BigQuery Python UDF. Let's try a `pandas`-like use case in which we want + # to apply a user defined function to every value in a `Series`, more + # specifically bucketize the `body_mass_g` value of the penguins, which is a + # real number, into a category, which is a string. + @bpd.udf( + dataset=your_bq_dataset_id, + name=your_bq_routine_id, + ) + def get_bucket(num: float) -> str: + if not num: + return "NA" + boundary = 4000 + return "at_or_above_4000" if num >= boundary else "below_4000" + + # Then we can apply the udf on the `Series` of interest via + # `apply` API and store the result in a new column in the DataFrame. + df = df.assign(body_mass_bucket=df["body_mass_g"].apply(get_bucket)) + + # This will add a new column `body_mass_bucket` in the DataFrame. You can + # preview the original value and the bucketized value side by side. + df[["body_mass_g", "body_mass_bucket"]].peek(10) + + # The above operation was possible by doing all the computation on the + # cloud through an underlying BigQuery Python UDF that was created to + # support the user's operations in the Python code. + + # The BigQuery Python UDF created to support the BigQuery DataFrames + # udf can be located via a property `bigframes_bigquery_function` + # set in the udf object. + print(f"Created BQ Python UDF: {get_bucket.bigframes_bigquery_function}") + + # If you have already defined a custom function in BigQuery, either via the + # BigQuery Google Cloud Console or with the `udf` decorator, + # or otherwise, you may use it with BigQuery DataFrames with the + # `read_gbq_function` method. More details are available via the `help` + # command. + help(bpd.read_gbq_function) + + existing_get_bucket_bq_udf = get_bucket.bigframes_bigquery_function + + # Here is an example of using `read_gbq_function` to load an existing + # BigQuery Python UDF. + df = bpd.read_gbq("bigquery-public-data.ml_datasets.penguins") + get_bucket_function = bpd.read_gbq_function(existing_get_bucket_bq_udf) + + df = df.assign(body_mass_bucket=df["body_mass_g"].apply(get_bucket_function)) + df.peek(10) + + # Let's continue trying other potential use cases of udf. Let's say we + # consider the `species`, `island` and `sex` of the penguins sensitive + # information and want to redact that by replacing with their hash code + # instead. Let's define another scalar custom function and decorate it + # as a udf. The custom function in this example has external package + # dependency, which can be specified via `packages` parameter. + @bpd.udf( + dataset=your_bq_dataset_id, + name=your_bq_routine_id, + packages=["cryptography"], + ) + def get_hash(input: str) -> str: + from cryptography.fernet import Fernet + + # handle missing value + if input is None: + input = "" + + key = Fernet.generate_key() + f = Fernet(key) + return f.encrypt(input.encode()).decode() + + # We can use this udf in another `pandas`-like API `map` that + # can be applied on a DataFrame + df_redacted = df[["species", "island", "sex"]].map(get_hash) + df_redacted.peek(10) + + # If the BigQuery routine is no longer needed, we can clean it up + # to free up any cloud quota + session = bpd.get_global_session() + session.bqclient.delete_routine(f"{your_bq_dataset_id}.{your_bq_routine_id}") + + # [END bigquery_dataframes_udf] diff --git a/samples/snippets/udf_test.py b/samples/snippets/udf_test.py new file mode 100644 index 00000000000..a352b4c8ce5 --- /dev/null +++ b/samples/snippets/udf_test.py @@ -0,0 +1,38 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + +import bigframes.pandas + +from . import udf + + +def test_udf_and_read_gbq_function( + capsys: pytest.CaptureFixture[str], + dataset_id: str, + routine_id: str, +) -> None: + # We need a fresh session since we're modifying connection options. + bigframes.pandas.close_session() + + # Determine project id, in this case prefer the one set in the environment + # variable GOOGLE_CLOUD_PROJECT (if any) + import os + + your_project_id = os.getenv("GOOGLE_CLOUD_PROJECT", "bigframes-dev") + + udf.run_udf_and_read_gbq_function(your_project_id, dataset_id, routine_id) + out, _ = capsys.readouterr() + assert "Created BQ Python UDF:" in out diff --git a/scripts/decrypt-secrets.sh b/scripts/decrypt-secrets.sh new file mode 100755 index 00000000000..120b0ddc436 --- /dev/null +++ b/scripts/decrypt-secrets.sh @@ -0,0 +1,46 @@ +#!/bin/bash + +# Copyright 2024 Google LLC All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )" +ROOT=$( dirname "$DIR" ) + +# Work from the project root. +cd $ROOT + +# Prevent it from overriding files. +# We recommend that sample authors use their own service account files and cloud project. +# In that case, they are supposed to prepare these files by themselves. +if [[ -f "testing/test-env.sh" ]] || \ + [[ -f "testing/service-account.json" ]] || \ + [[ -f "testing/client-secrets.json" ]]; then + echo "One or more target files exist, aborting." + exit 1 +fi + +# Use SECRET_MANAGER_PROJECT if set, fallback to cloud-devrel-kokoro-resources. +PROJECT_ID="${SECRET_MANAGER_PROJECT:-cloud-devrel-kokoro-resources}" + +gcloud secrets versions access latest --secret="python-docs-samples-test-env" \ + --project="${PROJECT_ID}" \ + > testing/test-env.sh +gcloud secrets versions access latest \ + --secret="python-docs-samples-service-account" \ + --project="${PROJECT_ID}" \ + > testing/service-account.json +gcloud secrets versions access latest \ + --secret="python-docs-samples-client-secrets" \ + --project="${PROJECT_ID}" \ + > testing/client-secrets.json diff --git a/scripts/readme-gen/readme_gen.py b/scripts/readme-gen/readme_gen.py new file mode 100644 index 00000000000..ceb1eada7c2 --- /dev/null +++ b/scripts/readme-gen/readme_gen.py @@ -0,0 +1,68 @@ +#!/usr/bin/env python + +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Generates READMEs using configuration defined in yaml.""" + +import argparse +import io +import os +import subprocess + +import jinja2 +import yaml + +jinja_env = jinja2.Environment( + trim_blocks=True, + loader=jinja2.FileSystemLoader( + os.path.abspath(os.path.join(os.path.dirname(__file__), "templates")) + ), + autoescape=True, +) + +README_TMPL = jinja_env.get_template("README.tmpl.rst") + + +def get_help(file): + return subprocess.check_output(["python", file, "--help"]).decode() + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("source") + parser.add_argument("--destination", default="README.rst") + + args = parser.parse_args() + + source = os.path.abspath(args.source) + root = os.path.dirname(source) + destination = os.path.join(root, args.destination) + + jinja_env.globals["get_help"] = get_help + + with io.open(source, "r") as f: + config = yaml.load(f) + + # This allows get_help to execute in the right directory. + os.chdir(root) + + output = README_TMPL.render(config) + + with io.open(destination, "w") as f: + f.write(output) + + +if __name__ == "__main__": + main() diff --git a/scripts/readme-gen/templates/README.tmpl.rst b/scripts/readme-gen/templates/README.tmpl.rst new file mode 100644 index 00000000000..4fd239765b0 --- /dev/null +++ b/scripts/readme-gen/templates/README.tmpl.rst @@ -0,0 +1,87 @@ +{# The following line is a lie. BUT! Once jinja2 is done with it, it will + become truth! #} +.. This file is automatically generated. Do not edit this file directly. + +{{product.name}} Python Samples +=============================================================================== + +.. image:: https://gstatic.com/cloudssh/images/open-btn.png + :target: https://console.cloud.google.com/cloudshell/open?git_repo=https://github.com/GoogleCloudPlatform/python-docs-samples&page=editor&open_in_editor={{folder}}/README.rst + + +This directory contains samples for {{product.name}}. {{product.description}} + +{{description}} + +.. _{{product.name}}: {{product.url}} + +{% if required_api_url %} +To run the sample, you need to enable the API at: {{required_api_url}} +{% endif %} + +{% if required_role %} +To run the sample, you need to have `{{required_role}}` role. +{% endif %} + +{{other_required_steps}} + +{% if setup %} +Setup +------------------------------------------------------------------------------- + +{% for section in setup %} + +{% include section + '.tmpl.rst' %} + +{% endfor %} +{% endif %} + +{% if samples %} +Samples +------------------------------------------------------------------------------- + +{% for sample in samples %} +{{sample.name}} ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ + +{% if not sample.hide_cloudshell_button %} +.. image:: https://gstatic.com/cloudssh/images/open-btn.png + :target: https://console.cloud.google.com/cloudshell/open?git_repo=https://github.com/GoogleCloudPlatform/python-docs-samples&page=editor&open_in_editor={{folder}}/{{sample.file}},{{folder}}/README.rst +{% endif %} + + +{{sample.description}} + +To run this sample: + +.. code-block:: bash + + $ python {{sample.file}} +{% if sample.show_help %} + + {{get_help(sample.file)|indent}} +{% endif %} + + +{% endfor %} +{% endif %} + +{% if cloud_client_library %} + +The client library +------------------------------------------------------------------------------- + +This sample uses the `Google Cloud Client Library for Python`_. +You can read the documentation for more details on API usage and use GitHub +to `browse the source`_ and `report issues`_. + +.. _Google Cloud Client Library for Python: + https://googlecloudplatform.github.io/google-cloud-python/ +.. _browse the source: + https://github.com/GoogleCloudPlatform/google-cloud-python +.. _report issues: + https://github.com/GoogleCloudPlatform/google-cloud-python/issues + +{% endif %} + +.. _Google Cloud SDK: https://cloud.google.com/sdk/ \ No newline at end of file diff --git a/scripts/readme-gen/templates/auth.tmpl.rst b/scripts/readme-gen/templates/auth.tmpl.rst new file mode 100644 index 00000000000..1446b94a5e3 --- /dev/null +++ b/scripts/readme-gen/templates/auth.tmpl.rst @@ -0,0 +1,9 @@ +Authentication +++++++++++++++ + +This sample requires you to have authentication setup. Refer to the +`Authentication Getting Started Guide`_ for instructions on setting up +credentials for applications. + +.. _Authentication Getting Started Guide: + https://cloud.google.com/docs/authentication/getting-started diff --git a/scripts/readme-gen/templates/auth_api_key.tmpl.rst b/scripts/readme-gen/templates/auth_api_key.tmpl.rst new file mode 100644 index 00000000000..11957ce2714 --- /dev/null +++ b/scripts/readme-gen/templates/auth_api_key.tmpl.rst @@ -0,0 +1,14 @@ +Authentication +++++++++++++++ + +Authentication for this service is done via an `API Key`_. To obtain an API +Key: + +1. Open the `Cloud Platform Console`_ +2. Make sure that billing is enabled for your project. +3. From the **Credentials** page, create a new **API Key** or use an existing + one for your project. + +.. _API Key: + https://developers.google.com/api-client-library/python/guide/aaa_apikeys +.. _Cloud Console: https://console.cloud.google.com/project?_ diff --git a/scripts/readme-gen/templates/install_deps.tmpl.rst b/scripts/readme-gen/templates/install_deps.tmpl.rst new file mode 100644 index 00000000000..6f069c6c87a --- /dev/null +++ b/scripts/readme-gen/templates/install_deps.tmpl.rst @@ -0,0 +1,29 @@ +Install Dependencies +++++++++++++++++++++ + +#. Clone python-docs-samples and change directory to the sample directory you want to use. + + .. code-block:: bash + + $ git clone https://github.com/GoogleCloudPlatform/python-docs-samples.git + +#. Install `pip`_ and `virtualenv`_ if you do not already have them. You may want to refer to the `Python Development Environment Setup Guide`_ for Google Cloud Platform for instructions. + + .. _Python Development Environment Setup Guide: + https://cloud.google.com/python/setup + +#. Create a virtualenv. Samples are compatible with Python 3.7+. + + .. code-block:: bash + + $ virtualenv env + $ source env/bin/activate + +#. Install the dependencies needed to run the samples. + + .. code-block:: bash + + $ pip install -r requirements.txt + +.. _pip: https://pip.pypa.io/ +.. _virtualenv: https://virtualenv.pypa.io/ diff --git a/scripts/readme-gen/templates/install_portaudio.tmpl.rst b/scripts/readme-gen/templates/install_portaudio.tmpl.rst new file mode 100644 index 00000000000..5ea33d18c00 --- /dev/null +++ b/scripts/readme-gen/templates/install_portaudio.tmpl.rst @@ -0,0 +1,35 @@ +Install PortAudio ++++++++++++++++++ + +Install `PortAudio`_. This is required by the `PyAudio`_ library to stream +audio from your computer's microphone. PyAudio depends on PortAudio for cross-platform compatibility, and is installed differently depending on the +platform. + +* For Mac OS X, you can use `Homebrew`_:: + + brew install portaudio + + **Note**: if you encounter an error when running `pip install` that indicates + it can't find `portaudio.h`, try running `pip install` with the following + flags:: + + pip install --global-option='build_ext' \ + --global-option='-I/usr/local/include' \ + --global-option='-L/usr/local/lib' \ + pyaudio + +* For Debian / Ubuntu Linux:: + + apt-get install portaudio19-dev python-all-dev + +* Windows may work without having to install PortAudio explicitly (it will get + installed with PyAudio). + +For more details, see the `PyAudio installation`_ page. + + +.. _PyAudio: https://people.csail.mit.edu/hubert/pyaudio/ +.. _PortAudio: http://www.portaudio.com/ +.. _PyAudio installation: + https://people.csail.mit.edu/hubert/pyaudio/#downloads +.. _Homebrew: http://brew.sh diff --git a/tests/js/package-lock.json b/tests/js/package-lock.json index aab813c3eb1..1d97bdfbff7 100644 --- a/tests/js/package-lock.json +++ b/tests/js/package-lock.json @@ -72,13 +72,13 @@ "license": "MIT" }, "node_modules/@babel/code-frame": { - "version": "7.29.0", - "resolved": "https://registry.npmjs.org/@babel/code-frame/-/code-frame-7.29.0.tgz", - "integrity": "sha512-9NhCeYjq9+3uxgdtp20LSiJXJvN0FeCtNGpJxuMFZ1Kv3cWUNb6DOhJwUvcVCzKGR66cw4njwM6hrJLqgOwbcw==", + "version": "7.29.7", + "resolved": "https://registry.npmjs.org/@babel/code-frame/-/code-frame-7.29.7.tgz", + "integrity": "sha512-Aup7aUOfpbAUg2ROOJN6Iw5f9DMBlzu0mIkm/malLQFN/YQgO48wCj0Kxa3sEHJvPVFg7siR+qRInwXd2qhQKw==", "dev": true, "license": "MIT", "dependencies": { - "@babel/helper-validator-identifier": "^7.28.5", + "@babel/helper-validator-identifier": "^7.29.7", "js-tokens": "^4.0.0", "picocolors": "^1.1.1" }, @@ -145,14 +145,14 @@ } }, "node_modules/@babel/generator": { - "version": "7.29.1", - "resolved": "https://registry.npmjs.org/@babel/generator/-/generator-7.29.1.tgz", - "integrity": "sha512-qsaF+9Qcm2Qv8SRIMMscAvG4O3lJ0F1GuMo5HR/Bp02LopNgnZBC/EkbevHFeGs4ls/oPz9v+Bsmzbkbe+0dUw==", + "version": "7.29.8", + "resolved": "https://registry.npmjs.org/@babel/generator/-/generator-7.29.8.tgz", + "integrity": "sha512-gZbepsdh3WDtgZKWL+vTPh71LSBrm/Y4/QDZBVCcYfmeTEEuoOYwlSy+G1StfJg+/Zy550u/3TATbm7qDbbMtg==", "dev": true, "license": "MIT", "dependencies": { - "@babel/parser": "^7.29.0", - "@babel/types": "^7.29.0", + "@babel/parser": "^7.29.8", + "@babel/types": "^7.29.8", "@jridgewell/gen-mapping": "^0.3.12", "@jridgewell/trace-mapping": "^0.3.28", "jsesc": "^3.0.2" @@ -289,9 +289,9 @@ } }, "node_modules/@babel/helper-globals": { - "version": "7.28.0", - "resolved": "https://registry.npmjs.org/@babel/helper-globals/-/helper-globals-7.28.0.tgz", - "integrity": "sha512-+W6cISkXFa1jXsDEdYA8HeevQT/FULhxzR99pxphltZcVaugps53THCeiWA8SguxxpSp3gKPiuYfSWopkLQ4hw==", + "version": "7.29.7", + "resolved": "https://registry.npmjs.org/@babel/helper-globals/-/helper-globals-7.29.7.tgz", + "integrity": "sha512-3nQVUAtvkKH9zahfWgw96Jc/uFOmjACE1kQz82E2lqWmHBgjzbNlsC22nuQTfahmWeQtTq5nQ/4Nnd2A1wj4zA==", "dev": true, "license": "MIT", "engines": { @@ -313,29 +313,29 @@ } }, "node_modules/@babel/helper-module-imports": { - "version": "7.28.6", - "resolved": "https://registry.npmjs.org/@babel/helper-module-imports/-/helper-module-imports-7.28.6.tgz", - "integrity": "sha512-l5XkZK7r7wa9LucGw9LwZyyCUscb4x37JWTPz7swwFE/0FMQAGpiWUZn8u9DzkSBWEcK25jmvubfpw2dnAMdbw==", + "version": "7.29.7", + "resolved": "https://registry.npmjs.org/@babel/helper-module-imports/-/helper-module-imports-7.29.7.tgz", + "integrity": "sha512-ejHwrQQYcm9xnTivShn2IDOlIzInN34AXskvq9QicvCtEzq1Vzclu/tKF8Jq1Cg8JG2GL6/EmjgsCT7lXepE3g==", "dev": true, "license": "MIT", "dependencies": { - "@babel/traverse": "^7.28.6", - "@babel/types": "^7.28.6" + "@babel/traverse": "^7.29.7", + "@babel/types": "^7.29.7" }, "engines": { "node": ">=6.9.0" } }, "node_modules/@babel/helper-module-transforms": { - "version": "7.28.6", - "resolved": "https://registry.npmjs.org/@babel/helper-module-transforms/-/helper-module-transforms-7.28.6.tgz", - "integrity": "sha512-67oXFAYr2cDLDVGLXTEABjdBJZ6drElUSI7WKp70NrpyISso3plG9SAGEF6y7zbha/wOzUByWWTJvEDVNIUGcA==", + "version": "7.29.7", + "resolved": "https://registry.npmjs.org/@babel/helper-module-transforms/-/helper-module-transforms-7.29.7.tgz", + "integrity": "sha512-UPUVSyXbOh627KiCIGQSgwWzGeBKLkaJ9PJEdrngIwMSzxLR4jS4+f1f1jb7VzBbg8nFLaYotvVPFCTqdrmTAg==", "dev": true, "license": "MIT", "dependencies": { - "@babel/helper-module-imports": "^7.28.6", - "@babel/helper-validator-identifier": "^7.28.5", - "@babel/traverse": "^7.28.6" + "@babel/helper-module-imports": "^7.29.7", + "@babel/helper-validator-identifier": "^7.29.7", + "@babel/traverse": "^7.29.7" }, "engines": { "node": ">=6.9.0" @@ -358,9 +358,9 @@ } }, "node_modules/@babel/helper-plugin-utils": { - "version": "7.28.6", - "resolved": "https://registry.npmjs.org/@babel/helper-plugin-utils/-/helper-plugin-utils-7.28.6.tgz", - "integrity": "sha512-S9gzZ/bz83GRysI7gAD4wPT/AI3uCnY+9xn+Mx/KPs2JwHJIz1W8PZkg2cqyt3RNOBM8ejcXhV6y8Og7ly/Dug==", + "version": "7.29.7", + "resolved": "https://registry.npmjs.org/@babel/helper-plugin-utils/-/helper-plugin-utils-7.29.7.tgz", + "integrity": "sha512-G7sHYigPY17oO5SYWnfD/0MTBwVR781S/JI643e/JhUYgVgWE/61SoW3NH9KWUKyKq5LVh3npif99Wkt6j86Jw==", "dev": true, "license": "MIT", "engines": { @@ -418,9 +418,9 @@ } }, "node_modules/@babel/helper-string-parser": { - "version": "7.27.1", - "resolved": "https://registry.npmjs.org/@babel/helper-string-parser/-/helper-string-parser-7.27.1.tgz", - "integrity": "sha512-qMlSxKbpRlAridDExk92nSobyDdpPijUq2DW6oDnUqd0iOGxmQjyqhMIihI9+zv4LPyZdRje2cavWPbCbWm3eA==", + "version": "7.29.7", + "resolved": "https://registry.npmjs.org/@babel/helper-string-parser/-/helper-string-parser-7.29.7.tgz", + "integrity": "sha512-Pb5ijPrZ89GDH8223L4UP8i6QApWxs04RbPQJTeWDV0/keR2E36MeKnyr6LYmUUvqRRI+Iv87SuF1W6ErINzYw==", "dev": true, "license": "MIT", "engines": { @@ -428,9 +428,9 @@ } }, "node_modules/@babel/helper-validator-identifier": { - "version": "7.28.5", - "resolved": "https://registry.npmjs.org/@babel/helper-validator-identifier/-/helper-validator-identifier-7.28.5.tgz", - "integrity": "sha512-qSs4ifwzKJSV39ucNjsvc6WVHs6b7S03sOh2OcHF9UHfVPqWWALUsNUVzhSBiItjRZoLHx7nIarVjqKVusUZ1Q==", + "version": "7.29.7", + "resolved": "https://registry.npmjs.org/@babel/helper-validator-identifier/-/helper-validator-identifier-7.29.7.tgz", + "integrity": "sha512-qehxGkRj55h/ff8EMaJ+cYhyaKlHIxqYDn682wQD7RNp9UujOQsHog2uS0r2vzr4pW+sXf90NeeayjcNaX3fFg==", "dev": true, "license": "MIT", "engines": { @@ -477,13 +477,13 @@ } }, "node_modules/@babel/parser": { - "version": "7.29.2", - "resolved": "https://registry.npmjs.org/@babel/parser/-/parser-7.29.2.tgz", - "integrity": "sha512-4GgRzy/+fsBa72/RZVJmGKPmZu9Byn8o4MoLpmNe1m8ZfYnz5emHLQz3U4gLud6Zwl0RZIcgiLD7Uq7ySFuDLA==", + "version": "7.29.8", + "resolved": "https://registry.npmjs.org/@babel/parser/-/parser-7.29.8.tgz", + "integrity": "sha512-E8lTAYNB1KW+FH+VGJuZM1ioAx2E6oVlvQFRrf5P8ZZmsiJXYAD9vTFV7yyEURNzgh1dFqMZuO6tUwcARbqFCA==", "dev": true, "license": "MIT", "dependencies": { - "@babel/types": "^7.29.0" + "@babel/types": "^7.29.8" }, "bin": { "parser": "bin/babel-parser.js" @@ -1300,16 +1300,16 @@ } }, "node_modules/@babel/plugin-transform-modules-systemjs": { - "version": "7.29.4", - "resolved": "https://registry.npmjs.org/@babel/plugin-transform-modules-systemjs/-/plugin-transform-modules-systemjs-7.29.4.tgz", - "integrity": "sha512-N7QmZ0xRZfjHOfZeQLJjwgX2zS9pdGHSVl/cjSGlo4dXMqvurfxXDMKY4RqEKzPozV78VMcd0lxyG13mlbKc4w==", + "version": "7.29.8", + "resolved": "https://registry.npmjs.org/@babel/plugin-transform-modules-systemjs/-/plugin-transform-modules-systemjs-7.29.8.tgz", + "integrity": "sha512-6iSnEK0zlkLKU4heofK/AdmRD4e2SHVpJMtrwnTCzhnaM98ria4rTrOXBBi45BTTYnJtO8txnPsX4fChYXkmeA==", "dev": true, "license": "MIT", "dependencies": { - "@babel/helper-module-transforms": "^7.28.6", - "@babel/helper-plugin-utils": "^7.28.6", - "@babel/helper-validator-identifier": "^7.28.5", - "@babel/traverse": "^7.29.0" + "@babel/helper-module-transforms": "^7.29.7", + "@babel/helper-plugin-utils": "^7.29.7", + "@babel/helper-validator-identifier": "^7.29.7", + "@babel/traverse": "^7.29.8" }, "engines": { "node": ">=6.9.0" @@ -1846,33 +1846,33 @@ } }, "node_modules/@babel/template": { - "version": "7.28.6", - "resolved": "https://registry.npmjs.org/@babel/template/-/template-7.28.6.tgz", - "integrity": "sha512-YA6Ma2KsCdGb+WC6UpBVFJGXL58MDA6oyONbjyF/+5sBgxY/dwkhLogbMT2GXXyU84/IhRw/2D1Os1B/giz+BQ==", + "version": "7.29.7", + "resolved": "https://registry.npmjs.org/@babel/template/-/template-7.29.7.tgz", + "integrity": "sha512-puq+Gf35oI24FeN11LkoUQFqv9uwNeWpxXZi/Ji3rRIoKAzKnxRaZ+Gkj0vKS9ZCiTESfng1N9LyOyXvo+m+Gg==", "dev": true, "license": "MIT", "dependencies": { - "@babel/code-frame": "^7.28.6", - "@babel/parser": "^7.28.6", - "@babel/types": "^7.28.6" + "@babel/code-frame": "^7.29.7", + "@babel/parser": "^7.29.7", + "@babel/types": "^7.29.7" }, "engines": { "node": ">=6.9.0" } }, "node_modules/@babel/traverse": { - "version": "7.29.0", - "resolved": "https://registry.npmjs.org/@babel/traverse/-/traverse-7.29.0.tgz", - "integrity": "sha512-4HPiQr0X7+waHfyXPZpWPfWL/J7dcN1mx9gL6WdQVMbPnF3+ZhSMs8tCxN7oHddJE9fhNE7+lxdnlyemKfJRuA==", + "version": "7.29.8", + "resolved": "https://registry.npmjs.org/@babel/traverse/-/traverse-7.29.8.tgz", + "integrity": "sha512-I5z7H3bf/41ktsNVLtpN0wAa336HkqIHQ5BuPLEhTkt1jVSyZpeNKIzTgEWmlxjdg81R0IgUCcaE+Ok3NvrfZg==", "dev": true, "license": "MIT", "dependencies": { - "@babel/code-frame": "^7.29.0", - "@babel/generator": "^7.29.0", - "@babel/helper-globals": "^7.28.0", - "@babel/parser": "^7.29.0", - "@babel/template": "^7.28.6", - "@babel/types": "^7.29.0", + "@babel/code-frame": "^7.29.7", + "@babel/generator": "^7.29.8", + "@babel/helper-globals": "^7.29.7", + "@babel/parser": "^7.29.8", + "@babel/template": "^7.29.7", + "@babel/types": "^7.29.8", "debug": "^4.3.1" }, "engines": { @@ -1880,14 +1880,14 @@ } }, "node_modules/@babel/types": { - "version": "7.29.0", - "resolved": "https://registry.npmjs.org/@babel/types/-/types-7.29.0.tgz", - "integrity": "sha512-LwdZHpScM4Qz8Xw2iKSzS+cfglZzJGvofQICy7W7v4caru4EaAmyUuO6BGrbyQ2mYV11W0U8j5mBhd14dd3B0A==", + "version": "7.29.8", + "resolved": "https://registry.npmjs.org/@babel/types/-/types-7.29.8.tgz", + "integrity": "sha512-Vj1jF3cPfxg7OAfoI7QnVKLoILlm2JF9pnVHrX8qx7AHMiYWT+NDAA7jChlNgRS4WTLc/fD1lXLmPixluj+3Gg==", "dev": true, "license": "MIT", "dependencies": { - "@babel/helper-string-parser": "^7.27.1", - "@babel/helper-validator-identifier": "^7.28.5" + "@babel/helper-string-parser": "^7.29.7", + "@babel/helper-validator-identifier": "^7.29.7" }, "engines": { "node": ">=6.9.0" @@ -3346,9 +3346,9 @@ } }, "node_modules/brace-expansion": { - "version": "2.1.0", - "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-2.1.0.tgz", - "integrity": "sha512-TN1kCZAgdgweJhWWpgKYrQaMNHcDULHkWwQIspdtjV4Y5aurRdZpjAqn6yX3FPqTA9ngHCc4hJxMAMgGfve85w==", + "version": "2.1.4", + "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-2.1.4.tgz", + "integrity": "sha512-hGfVzPxthbf3+2yjg/RBs60cB0FhqBS/zvdV/4wn4/BmN0bNMMHPc4V/BbFieqf1TKAGGAHnY4eSjajCl0f2Xg==", "dev": true, "license": "MIT", "dependencies": { @@ -6438,9 +6438,9 @@ } }, "node_modules/test-exclude/node_modules/brace-expansion": { - "version": "1.1.14", - "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-1.1.14.tgz", - "integrity": "sha512-MWPGfDxnyzKU7rNOW9SP/c50vi3xrmrua/+6hfPbCS2ABNWfx24vPidzvC7krjU/RTo235sV776ymlsMtGKj8g==", + "version": "1.1.18", + "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-1.1.18.tgz", + "integrity": "sha512-Edep/X9fGqVNmzKBVsDYIOtD+z1tuezV70LBjdCst9Tqu76lsnvRiZ6oTic1n+/BIwX6QDGAO94PN4N2SADvtw==", "dev": true, "license": "MIT", "dependencies": { @@ -6928,9 +6928,9 @@ } }, "node_modules/ws": { - "version": "8.18.3", - "resolved": "https://registry.npmjs.org/ws/-/ws-8.18.3.tgz", - "integrity": "sha512-PEIGCY5tSlUt50cqyMXfCzX+oOPqN0vuGqWzbcJ2xvnkzkq46oOpz7dQaTDBdfICb4N14+GARUDw2XV2N4tvzg==", + "version": "8.21.3", + "resolved": "https://registry.npmjs.org/ws/-/ws-8.21.3.tgz", + "integrity": "sha512-201TZ/kPWxoPr/OKWjquZR1SWKXcvxdH+e1xrx89b3YbmzLMFCLfnaG1HFIgWzJOEWZ7MvpK++odZufgYR50Rw==", "dev": true, "license": "MIT", "engines": { diff --git a/tests/system/large/bigquery/test_obj.py b/tests/system/large/bigquery/test_obj.py new file mode 100644 index 00000000000..dcca7580b14 --- /dev/null +++ b/tests/system/large/bigquery/test_obj.py @@ -0,0 +1,41 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + +import bigframes.bigquery as bbq + + +@pytest.fixture() +def objectrefs(bq_connection): + return bbq.obj.make_ref( + [ + "gs://cloud-samples-data/bigquery/tutorials/cymbal-pets/images/tick-terminator-for-dogs.png" + ], + bq_connection, + ) + + +def test_obj_fetch_metadata(objectrefs): + metadata = bbq.obj.fetch_metadata(objectrefs) + + result = metadata.to_pandas() + assert len(result) == len(objectrefs) + + +def test_obj_get_access_url(objectrefs): + access = bbq.obj.get_access_url(objectrefs, "r") + + result = access.to_pandas() + assert len(result) == len(objectrefs) diff --git a/tests/system/large/blob/test_function.py b/tests/system/large/blob/test_function.py new file mode 100644 index 00000000000..6c7d8121005 --- /dev/null +++ b/tests/system/large/blob/test_function.py @@ -0,0 +1,853 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging +import os +import traceback +from typing import Generator +import uuid + +from google.cloud import storage +import pandas as pd +import pytest + +import bigframes +from bigframes import dtypes +import bigframes.pandas as bpd + +pytest.skip("Skipping blob tests due to b/481790217", allow_module_level=True) + + +@pytest.fixture(scope="function") +def images_output_folder() -> Generator[str, None, None]: + id = uuid.uuid4().hex + folder = os.path.join("gs://bigframes_blob_test/output/", id) + yield folder + + # clean up + try: + cloud_storage_client = storage.Client() + bucket = cloud_storage_client.bucket("bigframes_blob_test") + blobs = bucket.list_blobs(prefix="output/" + id) + for blob in blobs: + blob.delete() + except Exception as exc: + traceback.print_exception(type(exc), exc, None) + + +@pytest.fixture(scope="function") +def images_output_uris(images_output_folder: str) -> list[str]: + return [ + os.path.join(images_output_folder, "img0.jpg"), + os.path.join(images_output_folder, "img1.jpg"), + ] + + +def test_blob_exif( + bq_connection: str, + session: bigframes.Session, +): + exif_image_df = session.from_glob_path( + "gs://bigframes_blob_test/images_exif/*", + name="blob_col", + connection=bq_connection, + ) + + actual = exif_image_df["blob_col"].blob.exif( + engine="pillow", connection=bq_connection, verbose=False + ) + expected = bpd.Series( + ['{"ExifOffset": 47, "Make": "MyCamera"}'], + session=session, + dtype=dtypes.JSON_DTYPE, + ) + pd.testing.assert_series_equal( + actual.to_pandas(), + expected.to_pandas(), + check_dtype=False, + check_index_type=False, + ) + + +def test_blob_exif_verbose( + bq_connection: str, + session: bigframes.Session, +): + exif_image_df = session.from_glob_path( + "gs://bigframes_blob_test/images_exif/*", + name="blob_col", + connection=bq_connection, + ) + + actual = exif_image_df["blob_col"].blob.exif( + engine="pillow", connection=bq_connection, verbose=True + ) + assert hasattr(actual, "struct") + actual_exploded = actual.struct.explode() + assert "status" in actual_exploded.columns + assert "content" in actual_exploded.columns + + status_series = actual_exploded["status"] + assert status_series.dtype == dtypes.STRING_DTYPE + + content_series = actual_exploded["content"] + assert content_series.dtype == dtypes.JSON_DTYPE + + +def test_blob_image_blur_to_series( + images_mm_df: bpd.DataFrame, + bq_connection: str, + images_output_uris: list[str], + session: bigframes.Session, +): + series = bpd.Series(images_output_uris, session=session).str.to_blob( + connection=bq_connection + ) + + actual = images_mm_df["blob_col"].blob.image_blur( + (8, 8), dst=series, connection=bq_connection, engine="opencv", verbose=False + ) + + expected_df = pd.DataFrame( + { + "uri": images_output_uris, + "version": [None, None], + "authorizer": [bq_connection.casefold(), bq_connection.casefold()], + "details": [None, None], + } + ) + pd.testing.assert_frame_equal( + actual.struct.explode().to_pandas(), + expected_df, + check_dtype=False, + check_index_type=False, + ) + + # verify the files exist + assert not actual.blob.size().isna().any() + + +def test_blob_image_blur_to_series_verbose( + images_mm_df: bpd.DataFrame, + bq_connection: str, + images_output_uris: list[str], + session: bigframes.Session, +): + series = bpd.Series(images_output_uris, session=session).str.to_blob( + connection=bq_connection + ) + + actual = images_mm_df["blob_col"].blob.image_blur( + (8, 8), dst=series, connection=bq_connection, engine="opencv", verbose=True + ) + + assert hasattr(actual, "struct") + actual_exploded = actual.struct.explode() + assert "status" in actual_exploded.columns + assert "content" in actual_exploded.columns + + status_series = actual_exploded["status"] + assert status_series.dtype == dtypes.STRING_DTYPE + + # Content should be blob objects for GCS destination + # verify the files exist + assert not actual.blob.size().isna().any() + + +def test_blob_image_blur_to_folder( + images_mm_df: bpd.DataFrame, + bq_connection: str, + images_output_folder: str, + images_output_uris: list[str], +): + actual = images_mm_df["blob_col"].blob.image_blur( + (8, 8), + dst=images_output_folder, + connection=bq_connection, + engine="opencv", + verbose=False, + ) + expected_df = pd.DataFrame( + { + "uri": images_output_uris, + "version": [None, None], + "authorizer": [bq_connection.casefold(), bq_connection.casefold()], + "details": [None, None], + } + ) + pd.testing.assert_frame_equal( + actual.struct.explode().to_pandas(), + expected_df, + check_dtype=False, + check_index_type=False, + ) + + # verify the files exist + assert not actual.blob.size().isna().any() + + +def test_blob_image_blur_to_folder_verbose( + images_mm_df: bpd.DataFrame, + bq_connection: str, + images_output_folder: str, + images_output_uris: list[str], +): + actual = images_mm_df["blob_col"].blob.image_blur( + (8, 8), + dst=images_output_folder, + connection=bq_connection, + engine="opencv", + verbose=True, + ) + assert hasattr(actual, "struct") + actual_exploded = actual.struct.explode() + assert "status" in actual_exploded.columns + assert "content" in actual_exploded.columns + + status_series = actual_exploded["status"] + assert status_series.dtype == dtypes.STRING_DTYPE + + content_series = actual_exploded["content"] + # Content should be blob objects for GCS destination + assert hasattr(content_series, "blob") + + # verify the files exist + assert not actual.blob.size().isna().any() + + +def test_blob_image_blur_to_bq(images_mm_df: bpd.DataFrame, bq_connection: str): + actual = images_mm_df["blob_col"].blob.image_blur( + (8, 8), connection=bq_connection, engine="opencv", verbose=False + ) + + assert isinstance(actual, bpd.Series) + assert len(actual) == 2 + assert actual.dtype == dtypes.BYTES_DTYPE + + +def test_blob_image_blur_to_bq_verbose(images_mm_df: bpd.DataFrame, bq_connection: str): + actual = images_mm_df["blob_col"].blob.image_blur( + (8, 8), connection=bq_connection, engine="opencv", verbose=True + ) + + assert isinstance(actual, bpd.Series) + assert len(actual) == 2 + + assert hasattr(actual, "struct") + actual_exploded = actual.struct.explode() + assert "status" in actual_exploded.columns + assert "content" in actual_exploded.columns + + status_series = actual_exploded["status"] + assert status_series.dtype == dtypes.STRING_DTYPE + + content_series = actual_exploded["content"] + assert content_series.dtype == dtypes.BYTES_DTYPE + + +def test_blob_image_resize_to_series( + images_mm_df: bpd.DataFrame, + bq_connection: str, + images_output_uris: list[str], + session: bigframes.Session, +): + series = bpd.Series(images_output_uris, session=session).str.to_blob( + connection=bq_connection + ) + + actual = images_mm_df["blob_col"].blob.image_resize( + (200, 300), + dst=series, + connection=bq_connection, + engine="opencv", + verbose=False, + ) + + expected_df = pd.DataFrame( + { + "uri": images_output_uris, + "version": [None, None], + "authorizer": [bq_connection.casefold(), bq_connection.casefold()], + "details": [None, None], + } + ) + pd.testing.assert_frame_equal( + actual.struct.explode().to_pandas(), + expected_df, + check_dtype=False, + check_index_type=False, + ) + + # verify the files exist + assert not actual.blob.size().isna().any() + + +def test_blob_image_resize_to_series_verbose( + images_mm_df: bpd.DataFrame, + bq_connection: str, + images_output_uris: list[str], + session: bigframes.Session, +): + series = bpd.Series(images_output_uris, session=session).str.to_blob( + connection=bq_connection + ) + + actual = images_mm_df["blob_col"].blob.image_resize( + (200, 300), + dst=series, + connection=bq_connection, + engine="opencv", + verbose=True, + ) + + assert hasattr(actual, "struct") + actual_exploded = actual.struct.explode() + assert "status" in actual_exploded.columns + assert "content" in actual_exploded.columns + + status_series = actual_exploded["status"] + assert status_series.dtype == dtypes.STRING_DTYPE + + content_series = actual_exploded["content"] + # Content should be blob objects for GCS destination + assert hasattr(content_series, "blob") + + # verify the files exist + assert not actual.blob.size().isna().any() + + +def test_blob_image_resize_to_folder( + images_mm_df: bpd.DataFrame, + bq_connection: str, + images_output_folder: str, + images_output_uris: list[str], +): + actual = images_mm_df["blob_col"].blob.image_resize( + (200, 300), + dst=images_output_folder, + connection=bq_connection, + engine="opencv", + verbose=False, + ) + + expected_df = pd.DataFrame( + { + "uri": images_output_uris, + "version": [None, None], + "authorizer": [bq_connection.casefold(), bq_connection.casefold()], + "details": [None, None], + } + ) + pd.testing.assert_frame_equal( + actual.struct.explode().to_pandas(), + expected_df, + check_dtype=False, + check_index_type=False, + ) + + # verify the files exist + assert not actual.blob.size().isna().any() + + +def test_blob_image_resize_to_folder_verbose( + images_mm_df: bpd.DataFrame, + bq_connection: str, + images_output_folder: str, + images_output_uris: list[str], +): + actual = images_mm_df["blob_col"].blob.image_resize( + (200, 300), + dst=images_output_folder, + connection=bq_connection, + engine="opencv", + verbose=True, + ) + + assert hasattr(actual, "struct") + actual_exploded = actual.struct.explode() + assert "status" in actual_exploded.columns + assert "content" in actual_exploded.columns + + status_series = actual_exploded["status"] + assert status_series.dtype == dtypes.STRING_DTYPE + + content_series = actual_exploded["content"] + # Content should be blob objects for GCS destination + assert hasattr(content_series, "blob") + + # verify the files exist + assert not content_series.blob.size().isna().any() + + +def test_blob_image_resize_to_bq(images_mm_df: bpd.DataFrame, bq_connection: str): + actual = images_mm_df["blob_col"].blob.image_resize( + (200, 300), connection=bq_connection, engine="opencv", verbose=False + ) + + assert isinstance(actual, bpd.Series) + assert len(actual) == 2 + assert actual.dtype == dtypes.BYTES_DTYPE + + +def test_blob_image_resize_to_bq_verbose( + images_mm_df: bpd.DataFrame, bq_connection: str +): + actual = images_mm_df["blob_col"].blob.image_resize( + (200, 300), connection=bq_connection, engine="opencv", verbose=True + ) + + assert isinstance(actual, bpd.Series) + assert len(actual) == 2 + + assert hasattr(actual, "struct") + actual_exploded = actual.struct.explode() + assert "status" in actual_exploded.columns + assert "content" in actual_exploded.columns + + status_series = actual_exploded["status"] + assert status_series.dtype == dtypes.STRING_DTYPE + + content_series = actual_exploded["content"] + assert content_series.dtype == dtypes.BYTES_DTYPE + + +def test_blob_image_normalize_to_series( + images_mm_df: bpd.DataFrame, + bq_connection: str, + images_output_uris: list[str], + session: bigframes.Session, +): + series = bpd.Series(images_output_uris, session=session).str.to_blob( + connection=bq_connection + ) + + actual = images_mm_df["blob_col"].blob.image_normalize( + alpha=50.0, + beta=150.0, + norm_type="minmax", + dst=series, + connection=bq_connection, + engine="opencv", + verbose=False, + ) + + expected_df = pd.DataFrame( + { + "uri": images_output_uris, + "version": [None, None], + "authorizer": [bq_connection.casefold(), bq_connection.casefold()], + "details": [None, None], + } + ) + pd.testing.assert_frame_equal( + actual.struct.explode().to_pandas(), + expected_df, + check_dtype=False, + check_index_type=False, + ) + + # verify the files exist + assert not actual.blob.size().isna().any() + + +def test_blob_image_normalize_to_series_verbose( + images_mm_df: bpd.DataFrame, + bq_connection: str, + images_output_uris: list[str], + session: bigframes.Session, +): + series = bpd.Series(images_output_uris, session=session).str.to_blob( + connection=bq_connection + ) + + actual = images_mm_df["blob_col"].blob.image_normalize( + alpha=50.0, + beta=150.0, + norm_type="minmax", + dst=series, + connection=bq_connection, + engine="opencv", + verbose=True, + ) + + assert hasattr(actual, "struct") + actual_exploded = actual.struct.explode() + assert "status" in actual_exploded.columns + assert "content" in actual_exploded.columns + + status_series = actual_exploded["status"] + assert status_series.dtype == dtypes.STRING_DTYPE + + content_series = actual_exploded["content"] + # Content should be blob objects for GCS destination + assert hasattr(content_series, "blob") + + +def test_blob_image_normalize_to_folder( + images_mm_df: bpd.DataFrame, + bq_connection: str, + images_output_folder: str, + images_output_uris: list[str], +): + actual = images_mm_df["blob_col"].blob.image_normalize( + alpha=50.0, + beta=150.0, + norm_type="minmax", + dst=images_output_folder, + connection=bq_connection, + engine="opencv", + verbose=False, + ) + + expected_df = pd.DataFrame( + { + "uri": images_output_uris, + "version": [None, None], + "authorizer": [bq_connection.casefold(), bq_connection.casefold()], + "details": [None, None], + } + ) + pd.testing.assert_frame_equal( + actual.struct.explode().to_pandas(), + expected_df, + check_dtype=False, + check_index_type=False, + ) + + # verify the files exist + assert not actual.blob.size().isna().any() + + +def test_blob_image_normalize_to_folder_verbose( + images_mm_df: bpd.DataFrame, + bq_connection: str, + images_output_folder: str, + images_output_uris: list[str], +): + actual = images_mm_df["blob_col"].blob.image_normalize( + alpha=50.0, + beta=150.0, + norm_type="minmax", + dst=images_output_folder, + connection=bq_connection, + engine="opencv", + verbose=True, + ) + + assert hasattr(actual, "struct") + actual_exploded = actual.struct.explode() + assert "status" in actual_exploded.columns + assert "content" in actual_exploded.columns + + status_series = actual_exploded["status"] + assert status_series.dtype == dtypes.STRING_DTYPE + + content_series = actual_exploded["content"] + # Content should be blob objects for GCS destination + assert hasattr(content_series, "blob") + + +def test_blob_image_normalize_to_bq(images_mm_df: bpd.DataFrame, bq_connection: str): + actual = images_mm_df["blob_col"].blob.image_normalize( + alpha=50.0, + beta=150.0, + norm_type="minmax", + connection=bq_connection, + engine="opencv", + verbose=False, + ) + + assert isinstance(actual, bpd.Series) + assert len(actual) == 2 + assert actual.dtype == dtypes.BYTES_DTYPE + + +def test_blob_image_normalize_to_bq_verbose( + images_mm_df: bpd.DataFrame, bq_connection: str +): + actual = images_mm_df["blob_col"].blob.image_normalize( + alpha=50.0, + beta=150.0, + norm_type="minmax", + connection=bq_connection, + engine="opencv", + verbose=True, + ) + + assert isinstance(actual, bpd.Series) + assert len(actual) == 2 + + assert hasattr(actual, "struct") + actual_exploded = actual.struct.explode() + assert "status" in actual_exploded.columns + assert "content" in actual_exploded.columns + + status_series = actual_exploded["status"] + assert status_series.dtype == dtypes.STRING_DTYPE + + content_series = actual_exploded["content"] + assert content_series.dtype == dtypes.BYTES_DTYPE + + +def test_blob_pdf_extract( + pdf_mm_df: bpd.DataFrame, + bq_connection: str, +): + actual = ( + pdf_mm_df["pdf"] + .blob.pdf_extract(connection=bq_connection, verbose=False, engine="pypdf") + .explode() + .to_pandas() + ) + + # check relative length + expected_text = "Sample PDF This is a testing file. Some dummy messages are used for testing purposes." + expected_len = len(expected_text) + + actual_text = actual[actual != ""].iloc[0] + actual_len = len(actual_text) + + relative_length_tolerance = 0.25 + min_acceptable_len = expected_len * (1 - relative_length_tolerance) + max_acceptable_len = expected_len * (1 + relative_length_tolerance) + assert min_acceptable_len <= actual_len <= max_acceptable_len, ( + f"Item (verbose=False): Extracted text length {actual_len} is outside the acceptable range " + f"[{min_acceptable_len:.0f}, {max_acceptable_len:.0f}]. " + f"Expected reference length was {expected_len}. " + ) + + # check for major keywords + major_keywords = ["Sample", "PDF", "testing", "dummy", "messages"] + for keyword in major_keywords: + assert ( + keyword.lower() in actual_text.lower() + ), f"Item (verbose=False): Expected keyword '{keyword}' not found in extracted text. " + + +def test_blob_pdf_extract_verbose( + pdf_mm_df: bpd.DataFrame, + bq_connection: str, +): + actual = ( + pdf_mm_df["pdf"] + .blob.pdf_extract(connection=bq_connection, verbose=True, engine="pypdf") + .explode() + .to_pandas() + ) + + # check relative length + expected_text = "Sample PDF This is a testing file. Some dummy messages are used for testing purposes." + expected_len = len(expected_text) + + # The first entry is for a file that doesn't exist, so we check the second one + successful_results = actual[actual.apply(lambda x: x["status"] == "")] + actual_text = successful_results.apply(lambda x: x["content"]).iloc[0] + actual_len = len(actual_text) + + relative_length_tolerance = 0.25 + min_acceptable_len = expected_len * (1 - relative_length_tolerance) + max_acceptable_len = expected_len * (1 + relative_length_tolerance) + assert min_acceptable_len <= actual_len <= max_acceptable_len, ( + f"Item (verbose=True): Extracted text length {actual_len} is outside the acceptable range " + f"[{min_acceptable_len:.0f}, {max_acceptable_len:.0f}]. " + f"Expected reference length was {expected_len}. " + ) + + # check for major keywords + major_keywords = ["Sample", "PDF", "testing", "dummy", "messages"] + for keyword in major_keywords: + assert ( + keyword.lower() in actual_text.lower() + ), f"Item (verbose=True): Expected keyword '{keyword}' not found in extracted text. " + + +def test_blob_pdf_chunk(pdf_mm_df: bpd.DataFrame, bq_connection: str): + actual = ( + pdf_mm_df["pdf"] + .blob.pdf_chunk( + connection=bq_connection, + chunk_size=50, + overlap_size=10, + verbose=False, + engine="pypdf", + ) + .explode() + .to_pandas() + ) + + # check relative length + expected_text = "Sample PDF This is a testing file. Some dummy messages are used for testing purposes." + expected_len = len(expected_text) + + # First entry is NA + actual_text = "".join(actual.dropna()) + actual_len = len(actual_text) + + relative_length_tolerance = 0.25 + min_acceptable_len = expected_len * (1 - relative_length_tolerance) + max_acceptable_len = expected_len * (1 + relative_length_tolerance) + assert min_acceptable_len <= actual_len <= max_acceptable_len, ( + f"Item (verbose=False): Extracted text length {actual_len} is outside the acceptable range " + f"[{min_acceptable_len:.0f}, {max_acceptable_len:.0f}]. " + f"Expected reference length was {expected_len}. " + ) + + # check for major keywords + major_keywords = ["Sample", "PDF", "testing", "dummy", "messages"] + for keyword in major_keywords: + assert ( + keyword.lower() in actual_text.lower() + ), f"Item (verbose=False): Expected keyword '{keyword}' not found in extracted text. " + + +def test_blob_pdf_chunk_verbose(pdf_mm_df: bpd.DataFrame, bq_connection: str): + actual = ( + pdf_mm_df["pdf"] + .blob.pdf_chunk( + connection=bq_connection, + chunk_size=50, + overlap_size=10, + verbose=True, + engine="pypdf", + ) + .explode() + .to_pandas() + ) + + # check relative length + expected_text = "Sample PDF This is a testing file. Some dummy messages are used for testing purposes." + expected_len = len(expected_text) + + # The first entry is for a file that doesn't exist, so we check the second one + successful_results = actual[actual.apply(lambda x: x["status"] == "")] + actual_text = "".join(successful_results.apply(lambda x: x["content"]).iloc[0]) + actual_len = len(actual_text) + + relative_length_tolerance = 0.25 + min_acceptable_len = expected_len * (1 - relative_length_tolerance) + max_acceptable_len = expected_len * (1 + relative_length_tolerance) + assert min_acceptable_len <= actual_len <= max_acceptable_len, ( + f"Item (verbose=True): Extracted text length {actual_len} is outside the acceptable range " + f"[{min_acceptable_len:.0f}, {max_acceptable_len:.0f}]. " + f"Expected reference length was {expected_len}. " + ) + + # check for major keywords + major_keywords = ["Sample", "PDF", "testing", "dummy", "messages"] + for keyword in major_keywords: + assert ( + keyword.lower() in actual_text.lower() + ), f"Item (verbose=True): Expected keyword '{keyword}' not found in extracted text. " + + +@pytest.mark.parametrize( + "model_name", + [ + "gemini-2.0-flash-001", + "gemini-2.0-flash-lite-001", + ], +) +def test_blob_transcribe( + audio_mm_df: bpd.DataFrame, + model_name: str, +): + actual = ( + audio_mm_df["audio"] + .blob.audio_transcribe( + model_name=model_name, # type: ignore + verbose=False, + ) + .to_pandas() + ) + + # check relative length + expected_text = "Now, as all books not primarily intended as picture-books consist principally of types composed to form letterpress" + expected_len = len(expected_text) + + actual_text = actual[0] + + if pd.isna(actual_text) or actual_text == "": + # Ensure the tests are robust to flakes in the model, which isn't + # particularly useful information for the bigframes team. + logging.warning(f"blob_transcribe() model {model_name} verbose=False failure") + return + + actual_len = len(actual_text) + + relative_length_tolerance = 0.2 + min_acceptable_len = expected_len * (1 - relative_length_tolerance) + max_acceptable_len = expected_len * (1 + relative_length_tolerance) + assert min_acceptable_len <= actual_len <= max_acceptable_len, ( + f"Item (verbose=False): Transcribed text length {actual_len} is outside the acceptable range " + f"[{min_acceptable_len:.0f}, {max_acceptable_len:.0f}]. " + f"Expected reference length was {expected_len}. " + ) + + # check for major keywords + major_keywords = ["book", "picture"] + for keyword in major_keywords: + assert ( + keyword.lower() in actual_text.lower() + ), f"Item (verbose=False): Expected keyword '{keyword}' not found in transcribed text. " + + +@pytest.mark.parametrize( + "model_name", + [ + "gemini-2.0-flash-001", + "gemini-2.0-flash-lite-001", + ], +) +def test_blob_transcribe_verbose( + audio_mm_df: bpd.DataFrame, + model_name: str, +): + actual = ( + audio_mm_df["audio"] + .blob.audio_transcribe( + model_name=model_name, # type: ignore + verbose=True, + ) + .to_pandas() + ) + + # check relative length + expected_text = "Now, as all books not primarily intended as picture-books consist principally of types composed to form letterpress" + expected_len = len(expected_text) + + actual_text = actual[0]["content"] + + if pd.isna(actual_text) or actual_text == "": + # Ensure the tests are robust to flakes in the model, which isn't + # particularly useful information for the bigframes team. + logging.warning(f"blob_transcribe() model {model_name} verbose=True failure") + return + + actual_len = len(actual_text) + + relative_length_tolerance = 0.2 + min_acceptable_len = expected_len * (1 - relative_length_tolerance) + max_acceptable_len = expected_len * (1 + relative_length_tolerance) + assert min_acceptable_len <= actual_len <= max_acceptable_len, ( + f"Item (verbose=True): Transcribed text length {actual_len} is outside the acceptable range " + f"[{min_acceptable_len:.0f}, {max_acceptable_len:.0f}]. " + f"Expected reference length was {expected_len}. " + ) + + # check for major keywords + major_keywords = ["book", "picture"] + for keyword in major_keywords: + assert ( + keyword.lower() in actual_text.lower() + ), f"Item (verbose=True): Expected keyword '{keyword}' not found in transcribed text. " diff --git a/tests/system/large/ml/test_multimodal_llm.py b/tests/system/large/ml/test_multimodal_llm.py new file mode 100644 index 00000000000..f94f0f1dee6 --- /dev/null +++ b/tests/system/large/ml/test_multimodal_llm.py @@ -0,0 +1,106 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pandas as pd +import pyarrow as pa +import pytest + +from bigframes.ml import llm +import bigframes.pandas as bpd +from bigframes.testing import utils + + +@pytest.mark.parametrize( + "model_name", + ( + "gemini-2.0-flash-001", + "gemini-2.0-flash-lite-001", + ), +) +@pytest.mark.flaky(retries=2) +def test_gemini_text_generator_multimodal_input( + images_mm_df: bpd.DataFrame, model_name, session, bq_connection +): + gemini_text_generator_model = llm.GeminiTextGenerator( + model_name=model_name, connection_name=bq_connection, session=session + ) + pd_df = gemini_text_generator_model.predict( + images_mm_df, prompt=["Describe", images_mm_df["blob_col"]] + ).to_pandas() + utils.check_pandas_df_schema_and_index( + pd_df, + columns=utils.ML_GENERATE_TEXT_OUTPUT + ["blob_col"], + index=2, + col_exact=False, + ) + + +@pytest.mark.flaky(retries=2) +def test_multimodal_embedding_generator_predict_default_params_success( + images_mm_df, session, bq_connection +): + text_embedding_model = llm.MultimodalEmbeddingGenerator( + connection_name=bq_connection, session=session + ) + df = text_embedding_model.predict(images_mm_df).to_pandas() + utils.check_pandas_df_schema_and_index( + df, + columns=utils.ML_MULTIMODAL_GENERATE_EMBEDDING_OUTPUT, + index=2, + col_exact=False, + ) + assert len(df["ml_generate_embedding_result"][0]) == 1408 + + +@pytest.mark.parametrize( + "model_name", + ("gemini-2.0-flash-001",), +) +@pytest.mark.flaky(retries=2) +def test_gemini_text_generator_multimodal_structured_output( + images_mm_df: bpd.DataFrame, model_name, session, bq_connection +): + gemini_text_generator_model = llm.GeminiTextGenerator( + model_name=model_name, connection_name=bq_connection, session=session + ) + output_schema = { + "bool_output": "bool", + "int_output": "int64", + "float_output": "float64", + "str_output": "string", + "array_output": "array", + "struct_output": "struct", + } + df = gemini_text_generator_model.predict( + images_mm_df, + prompt=["Describe", images_mm_df["blob_col"]], + output_schema=output_schema, + ) + assert df["bool_output"].dtype == pd.BooleanDtype() + assert df["int_output"].dtype == pd.Int64Dtype() + assert df["float_output"].dtype == pd.Float64Dtype() + assert df["str_output"].dtype == pd.StringDtype(storage="pyarrow") + assert df["array_output"].dtype == pd.ArrowDtype(pa.list_(pa.int64())) + assert df["struct_output"].dtype == pd.ArrowDtype( + pa.struct([("number", pa.int64())]) + ) + + pd_df = df.to_pandas() + utils.check_pandas_df_schema_and_index( + pd_df, + columns=list(output_schema.keys()) + + ["blob_col", "prompt", "full_response", "status"], + index=2, + col_exact=False, + ) diff --git a/tests/system/small/blob/test_io.py b/tests/system/small/blob/test_io.py new file mode 100644 index 00000000000..c89fb4c6e6e --- /dev/null +++ b/tests/system/small/blob/test_io.py @@ -0,0 +1,132 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +import pandas as pd +import pytest + +import bigframes +import bigframes.pandas as bpd + +pytest.skip("Skipping blob tests due to b/481790217", allow_module_level=True) + + +idisplay = pytest.importorskip("IPython.display") + + +def test_blob_create_from_uri_str( + bq_connection: str, session: bigframes.Session, images_uris +): + uri_series = bpd.Series(images_uris, session=session) + blob_series = uri_series.str.to_blob(connection=bq_connection) + + pd_blob_df = blob_series.struct.explode().to_pandas() + expected_pd_df = pd.DataFrame( + { + "uri": images_uris, + "version": [None, None], + "authorizer": [bq_connection.casefold(), bq_connection.casefold()], + "details": [None, None], + } + ) + + pd.testing.assert_frame_equal( + pd_blob_df, expected_pd_df, check_dtype=False, check_index_type=False + ) + + +def test_blob_create_from_glob_path( + bq_connection: str, session: bigframes.Session, images_gcs_path, images_uris +): + blob_df = session.from_glob_path( + images_gcs_path, connection=bq_connection, name="blob_col" + ) + pd_blob_df = ( + blob_df["blob_col"] + .struct.explode() + .to_pandas() + .sort_values("uri") + .reset_index(drop=True) + ) + + expected_df = pd.DataFrame( + { + "uri": images_uris, + "version": [None, None], + "authorizer": [bq_connection.casefold(), bq_connection.casefold()], + "details": [None, None], + } + ) + + pd.testing.assert_frame_equal( + pd_blob_df, expected_df, check_dtype=False, check_index_type=False + ) + + +def test_blob_create_read_gbq_object_table( + bq_connection: str, session: bigframes.Session, images_gcs_path, images_uris +): + obj_table = session._create_object_table(images_gcs_path, bq_connection) + + blob_df = session.read_gbq_object_table(obj_table, name="blob_col") + pd_blob_df = ( + blob_df["blob_col"] + .struct.explode() + .to_pandas() + .sort_values("uri") + .reset_index(drop=True) + ) + expected_df = pd.DataFrame( + { + "uri": images_uris, + "version": [None, None], + "authorizer": [bq_connection.casefold(), bq_connection.casefold()], + "details": [None, None], + } + ) + + pd.testing.assert_frame_equal( + pd_blob_df, expected_df, check_dtype=False, check_index_type=False + ) + + +def test_display_images(monkeypatch, images_mm_df: bpd.DataFrame): + mock_display = mock.Mock() + monkeypatch.setattr(idisplay, "display", mock_display) + + images_mm_df["blob_col"].blob.display() + + for call in mock_display.call_args_list: + args, _ = call + arg = args[0] + assert isinstance(arg, idisplay.Image) + + +def test_display_nulls( + monkeypatch, + bq_connection: str, + session: bigframes.Session, +): + uri_series = bpd.Series([None, None, None], dtype="string", session=session) + blob_series = uri_series.str.to_blob(connection=bq_connection) + mock_display = mock.Mock() + monkeypatch.setattr(idisplay, "display", mock_display) + + blob_series.blob.display() + + for call in mock_display.call_args_list: + args, _ = call + arg = args[0] + assert arg == "" diff --git a/tests/system/small/blob/test_urls.py b/tests/system/small/blob/test_urls.py new file mode 100644 index 00000000000..b2dd6604343 --- /dev/null +++ b/tests/system/small/blob/test_urls.py @@ -0,0 +1,31 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + +import bigframes.pandas as bpd + +pytest.skip("Skipping blob tests due to b/481790217", allow_module_level=True) + + +def test_blob_read_url(images_mm_df: bpd.DataFrame): + urls = images_mm_df["blob_col"].blob.read_url() + + assert urls.str.startswith("https://storage.googleapis.com/").all() + + +def test_blob_write_url(images_mm_df: bpd.DataFrame): + urls = images_mm_df["blob_col"].blob.write_url() + + assert urls.str.startswith("https://storage.googleapis.com/").all() diff --git a/tests/unit/bigquery/_operations/test_io.py b/tests/unit/bigquery/_operations/test_io.py new file mode 100644 index 00000000000..b5dc9544aa8 --- /dev/null +++ b/tests/unit/bigquery/_operations/test_io.py @@ -0,0 +1,40 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +import pytest + +import bigframes.bigquery._operations.io +import bigframes.session + + +@pytest.fixture +def mock_session(): + return mock.create_autospec(spec=bigframes.session.Session) + + +@mock.patch("bigframes.bigquery._operations.io._get_table_metadata") +def test_load_data(get_table_metadata_mock, mock_session): + bigframes.bigquery._operations.io.load_data( + "my-project.my_dataset.my_table", + columns={"col1": "INT64", "col2": "STRING"}, + from_files_options={"format": "CSV", "uris": ["gs://bucket/path*"]}, + session=mock_session, + ) + mock_session.read_gbq_query.assert_called_once() + generated_sql = mock_session.read_gbq_query.call_args[0][0] + expected = "LOAD DATA INTO `my-project.my_dataset.my_table` (\n `col1` INT64,\n `col2` STRING\n) FROM FILES (format='CSV', uris=['gs://bucket/path*'])" + assert generated_sql == expected + get_table_metadata_mock.assert_called_once() diff --git a/tests/unit/core/compile/sqlglot/conftest.py b/tests/unit/core/compile/sqlglot/conftest.py new file mode 100644 index 00000000000..cb5a14b690e --- /dev/null +++ b/tests/unit/core/compile/sqlglot/conftest.py @@ -0,0 +1,280 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pathlib +import typing + +from google.cloud import bigquery +import pandas as pd +import pyarrow as pa +import pytest + +from bigframes import dtypes +import bigframes.core as core +import bigframes.pandas as bpd +import bigframes.testing.mocks as mocks +import bigframes.testing.utils + +CURRENT_DIR = pathlib.Path(__file__).parent +DATA_DIR = CURRENT_DIR.parent.parent.parent.parent / "data" + + +def _create_compiler_session(table_name, table_schema): + """Helper function to create a compiler session.""" + from bigframes.testing import compiler_session + + anonymous_dataset = bigquery.DatasetReference.from_string( + "bigframes-dev.sqlglot_test" + ) + session = mocks.create_bigquery_session( + table_name=table_name, + table_schema=table_schema, + anonymous_dataset=anonymous_dataset, + ) + session._executor = compiler_session.SQLCompilerExecutor() + return session + + +@pytest.fixture(scope="session") +def compiler_session(scalar_types_table_schema): + """Compiler session for scalar types.""" + return _create_compiler_session("scalar_types", scalar_types_table_schema) + + +@pytest.fixture(scope="session") +def compiler_session_w_repeated_types(repeated_types_table_schema): + """Compiler session for repeated data types.""" + return _create_compiler_session("repeated_types", repeated_types_table_schema) + + +@pytest.fixture(scope="session") +def compiler_session_w_nested_structs_types(nested_structs_types_table_schema): + """Compiler session for nested STRUCT data types.""" + return _create_compiler_session( + "nested_structs_types", nested_structs_types_table_schema + ) + + +@pytest.fixture(scope="session") +def compiler_session_w_json_types(json_types_table_schema): + """Compiler session for JSON data types.""" + return _create_compiler_session("json_types", json_types_table_schema) + + +@pytest.fixture(scope="session") +def scalar_types_table_schema() -> typing.Sequence[bigquery.SchemaField]: + return [ + bigquery.SchemaField("bool_col", "BOOLEAN"), + bigquery.SchemaField("bytes_col", "BYTES"), + bigquery.SchemaField("date_col", "DATE"), + bigquery.SchemaField("datetime_col", "DATETIME"), + bigquery.SchemaField("geography_col", "GEOGRAPHY"), + bigquery.SchemaField("int64_col", "INTEGER"), + bigquery.SchemaField("int64_too", "INTEGER"), + bigquery.SchemaField("numeric_col", "NUMERIC"), + bigquery.SchemaField("float64_col", "FLOAT"), + bigquery.SchemaField("rowindex", "INTEGER"), + bigquery.SchemaField("rowindex_2", "INTEGER", mode="REQUIRED"), + bigquery.SchemaField("string_col", "STRING"), + bigquery.SchemaField("time_col", "TIME"), + bigquery.SchemaField("timestamp_col", "TIMESTAMP"), + bigquery.SchemaField("duration_col", "INTEGER"), + ] + + +@pytest.fixture(scope="session") +def scalar_types_df(compiler_session) -> bpd.DataFrame: + """Returns a BigFrames DataFrame containing all scalar types and using the `rowindex` + column as the index.""" + bf_df = compiler_session._loader.read_gbq_table( + "bigframes-dev.sqlglot_test.scalar_types", + enable_snapshot=False, + ) + bf_df = bf_df.set_index("rowindex", drop=False) + return bf_df + + +@pytest.fixture(scope="session") +def scalar_types_pandas_df() -> pd.DataFrame: + """Returns a pandas DataFrame containing all scalar types and using the `rowindex` + column as the index.""" + # TODO: add tests for empty dataframes + df = pd.read_json( + DATA_DIR / "scalars.jsonl", + lines=True, + ) + bigframes.testing.utils.convert_pandas_dtypes(df, bytes_col=True) + + df = df.set_index("rowindex", drop=False) + return df + + +@pytest.fixture(scope="module") +def scalar_types_array_value( + scalar_types_pandas_df: pd.DataFrame, compiler_session: bigframes.Session +) -> core.ArrayValue: + managed_data_source = core.local_data.ManagedArrowTable.from_pandas( + scalar_types_pandas_df + ) + return core.ArrayValue.from_managed(managed_data_source, compiler_session) + + +@pytest.fixture(scope="session") +def nested_structs_types_table_schema() -> typing.Sequence[bigquery.SchemaField]: + return [ + bigquery.SchemaField("id", "INTEGER"), + bigquery.SchemaField( + "people", + "RECORD", + fields=[ + bigquery.SchemaField("name", "STRING"), + bigquery.SchemaField("age", "INTEGER"), + bigquery.SchemaField( + "address", + "RECORD", + fields=[ + bigquery.SchemaField("city", "STRING"), + bigquery.SchemaField("country", "STRING"), + ], + ), + ], + ), + ] + + +@pytest.fixture(scope="session") +def nested_structs_types_df(compiler_session_w_nested_structs_types) -> bpd.DataFrame: + """Returns a BigFrames DataFrame containing all scalar types and using the `rowindex` + column as the index.""" + bf_df = compiler_session_w_nested_structs_types._loader.read_gbq_table( + "bigframes-dev.sqlglot_test.nested_structs_types", + enable_snapshot=False, + ) + bf_df = bf_df.set_index("id", drop=False) + return bf_df + + +@pytest.fixture(scope="session") +def nested_structs_pandas_df() -> pd.DataFrame: + """Returns a pandas DataFrame containing STRUCT types and using the `id` + column as the index.""" + + df = pd.read_json( + DATA_DIR / "nested_structs.jsonl", + lines=True, + ) + df = df.set_index("id") + + address_struct_schema = pa.struct( + [pa.field("city", pa.string()), pa.field("country", pa.string())] + ) + person_struct_schema = pa.struct( + [ + pa.field("name", pa.string()), + pa.field("age", pa.int64()), + pa.field("address", address_struct_schema), + ] + ) + df["person"] = df["person"].astype(pd.ArrowDtype(person_struct_schema)) + return df + + +@pytest.fixture(scope="session") +def repeated_types_table_schema() -> typing.Sequence[bigquery.SchemaField]: + return [ + bigquery.SchemaField("rowindex", "INTEGER"), + bigquery.SchemaField("int_list_col", "INTEGER", "REPEATED"), + bigquery.SchemaField("bool_list_col", "BOOLEAN", "REPEATED"), + bigquery.SchemaField("float_list_col", "FLOAT", "REPEATED"), + bigquery.SchemaField("date_list_col", "DATE", "REPEATED"), + bigquery.SchemaField("date_time_list_col", "DATETIME", "REPEATED"), + bigquery.SchemaField("numeric_list_col", "NUMERIC", "REPEATED"), + bigquery.SchemaField("string_list_col", "STRING", "REPEATED"), + ] + + +@pytest.fixture(scope="session") +def repeated_types_df(compiler_session_w_repeated_types) -> bpd.DataFrame: + """Returns a BigFrames DataFrame containing all scalar types and using the `rowindex` + column as the index.""" + bf_df = compiler_session_w_repeated_types._loader.read_gbq_table( + "bigframes-dev.sqlglot_test.repeated_types", + enable_snapshot=False, + ) + bf_df = bf_df.set_index("rowindex", drop=False) + return bf_df + + +@pytest.fixture(scope="session") +def repeated_types_pandas_df() -> pd.DataFrame: + """Returns a pandas DataFrame containing LIST types and using the `rowindex` + column as the index.""" + + df = pd.read_json( + DATA_DIR / "repeated.jsonl", + lines=True, + ) + # TODO: add dtype conversion here if needed. + df = df.set_index("rowindex") + return df + + +@pytest.fixture(scope="session") +def json_types_table_schema() -> typing.Sequence[bigquery.SchemaField]: + return [ + bigquery.SchemaField("rowindex", "INTEGER"), + bigquery.SchemaField("json_col", "JSON"), + ] + + +@pytest.fixture(scope="session") +def json_types_df(compiler_session_w_json_types) -> bpd.DataFrame: + """Returns a BigFrames DataFrame containing JSON types and using the `rowindex` + column as the index.""" + bf_df = compiler_session_w_json_types._loader.read_gbq_table( + "bigframes-dev.sqlglot_test.json_types", + enable_snapshot=False, + ) + # TODO(b/427305807): Why `drop=False` will produce two "rowindex" columns? + bf_df = bf_df.set_index("rowindex", drop=True) + return bf_df + + +@pytest.fixture(scope="session") +def json_pandas_df() -> pd.DataFrame: + """Returns a pandas DataFrame containing JSON types and using the `rowindex` + column as the index.""" + json_data = [ + "null", + "true", + "100", + "0.98", + '"a string"', + "[]", + "[1, 2, 3]", + '[{"a": 1}, {"a": 2}, {"a": null}, {}]', + '"100"', + '{"date": "2024-07-16"}', + '{"int_value": 2, "null_filed": null}', + '{"list_data": [10, 20, 30]}', + ] + df = pd.DataFrame( + { + "rowindex": pd.Series(range(len(json_data)), dtype=dtypes.INT_DTYPE), + "json_col": pd.Series(json_data, dtype=dtypes.JSON_DTYPE), + }, + ) + # TODO(b/427305807): Why `drop=False` will produce two "rowindex" columns? + df = df.set_index("rowindex", drop=True) + return df diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_blob_ops/test_obj_fetch_metadata/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_blob_ops/test_obj_fetch_metadata/out.sql new file mode 100644 index 00000000000..ca6f5842df1 --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_blob_ops/test_obj_fetch_metadata/out.sql @@ -0,0 +1,6 @@ +SELECT + `rowindex`, + OBJ.FETCH_METADATA( + OBJ.MAKE_REF(`string_col`, 'bigframes-dev.test-region.bigframes-default-connection') + ).`version` +FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_blob_ops/test_obj_get_access_url/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_blob_ops/test_obj_get_access_url/out.sql new file mode 100644 index 00000000000..6d612ee6b87 --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_blob_ops/test_obj_get_access_url/out.sql @@ -0,0 +1,10 @@ +SELECT + `rowindex`, + JSON_VALUE( + OBJ.GET_ACCESS_URL( + OBJ.MAKE_REF(`string_col`, 'bigframes-dev.test-region.bigframes-default-connection'), + 'R' + ), + '$.access_urls.read_url' + ) AS `string_col` +FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_blob_ops/test_obj_get_access_url_with_duration/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_blob_ops/test_obj_get_access_url_with_duration/out.sql new file mode 100644 index 00000000000..2e8b60230fa --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_blob_ops/test_obj_get_access_url_with_duration/out.sql @@ -0,0 +1,3 @@ +SELECT + OBJ.GET_ACCESS_URL(`string_col`, 'READ', INTERVAL 3600 MICROSECOND) AS `string_col` +FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_blob_ops/test_obj_make_ref/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_blob_ops/test_obj_make_ref/out.sql new file mode 100644 index 00000000000..74ca601cd5d --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_blob_ops/test_obj_make_ref/out.sql @@ -0,0 +1,4 @@ +SELECT + `rowindex`, + OBJ.MAKE_REF(`string_col`, 'bigframes-dev.test-region.bigframes-default-connection') AS `string_col` +FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_blob_ops/test_obj_make_ref_json/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_blob_ops/test_obj_make_ref_json/out.sql new file mode 100644 index 00000000000..dc84b3bec12 --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_blob_ops/test_obj_make_ref_json/out.sql @@ -0,0 +1,3 @@ +SELECT + OBJ.MAKE_REF(`string_col`) AS `string_col` +FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/test_blob_ops.py b/tests/unit/core/compile/sqlglot/expressions/test_blob_ops.py new file mode 100644 index 00000000000..ac032f46e68 --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/test_blob_ops.py @@ -0,0 +1,60 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + +from bigframes import operations as ops +import bigframes.pandas as bpd +from bigframes.testing import utils + +pytest.importorskip("pytest_snapshot") + + +def test_obj_fetch_metadata(scalar_types_df: bpd.DataFrame, snapshot): + blob_s = scalar_types_df["string_col"].str.to_blob() + sql = blob_s.blob.version().to_frame().sql + snapshot.assert_match(sql, "out.sql") + + +def test_obj_get_access_url(scalar_types_df: bpd.DataFrame, snapshot): + blob_s = scalar_types_df["string_col"].str.to_blob() + sql = blob_s.blob.read_url().to_frame().sql + snapshot.assert_match(sql, "out.sql") + + +def test_obj_get_access_url_with_duration(scalar_types_df: bpd.DataFrame, snapshot): + col_name = "string_col" + bf_df = scalar_types_df[[col_name]] + sql = utils._apply_ops_to_sql( + bf_df, + [ops.ObjGetAccessUrl(mode="READ", duration=3600).as_expr(col_name)], + [col_name], + ) + snapshot.assert_match(sql, "out.sql") + + +def test_obj_make_ref(scalar_types_df: bpd.DataFrame, snapshot): + blob_df = scalar_types_df["string_col"].str.to_blob() + snapshot.assert_match(blob_df.to_frame().sql, "out.sql") + + +def test_obj_make_ref_json(scalar_types_df: bpd.DataFrame, snapshot): + col_name = "string_col" + bf_df = scalar_types_df[[col_name]] + sql = utils._apply_ops_to_sql( + bf_df, + [ops.obj_make_ref_json_op.as_expr(col_name)], + [col_name], + ) + snapshot.assert_match(sql, "out.sql")